dylanyunlong/project_6
1
0
Fork 0
Code Issues Pull Requests Actions Projects Releases Wiki Activity

ref(EX): import upstream ILU kernels + xllm MoE CUDA sources into ex_engine

Browse Source 
Copied from upstream_ref (NOT rewritten — exact upstream code):

ixformer C++ API (the authoritative header):
  include/ixformer.h — ixformer::infer namespace: topk_softmax,
    moe_compute_token_index_api, moe_w16a16_group_gemm, moe_expand_input,
    moe_output_reduce_sum, silu_and_mul, rms_norm, xllm_paged_attention, etc.
  include/ilu_ops_api.h — xllm::kernel::ilu namespace: moe_active_topk,
    moe_gen_idx, moe_expand_input, group_gemm, moe_combine_result,
    batch_prefill, batch_decode, rms_norm, matmul, act_and_mul, etc.

ILU kernel wrappers (call ixformer::infer directly):
  csrc/ilu_kernel_fused_moe.cpp — topk routing + gen_idx + expand + combine
  csrc/ilu_kernel_group_gemm.cpp — batched expert GEMM
  csrc/ilu_kernel_{activation,norm,rope,matmul,attention}.cpp

ILU layer implementations (full pipeline):
  csrc/ilu_layer_fused_moe.{cpp,h} — 797 lines, the complete MoE pipeline
    that competitor 168 ran as corex_moe.py
  csrc/ilu_layer_attention.{cpp,h} — prefill/decode attention dispatch

CUDA MoE kernels (from xllm + ds_vllm):
  csrc/moe/moe_topk_softmax_kernels.cuh — CUB BlockReduce + warp topk
  csrc/moe/moe_topk_sigmoid_kernels.cuh — sigmoid scoring variant
  csrc/moe/moe_topk.cuh + moe_fused_topk.cu — entry points
  csrc/moe/moeTopKFuncs.cuh — TRT-LLM derived vllm-compatible topk
  csrc/moe/moe_ops.h + moe_align_sum_kernels.cu — alignment kernels

Common layer headers:
  csrc/common_fused_moe{,_base}.h + common_moe_fused_topk.{cpp,h}
This commit is contained in:
project6-dev
2026-08-10 03:59:37 +00:00
parent dba027fded
commit f4e2264a83
26 changed files with 4522 additions and 16 deletions
Download Patch File Download Diff File Expand all files Collapse all files

123
ex_engine/csrc/moe/fused_moe_cuda.cpp Normal file
View File

@@ -0,0 +1,123 @@
/* Copyright 2026 The xLLM Authors. All Rights Reserved.
Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at
https://github.com/jd-opensource/xllm/blob/main/LICENSE
Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
==============================================================================*/
#include "kernels/cuda/cuda_ops_api.h"
#include "kernels/cuda/utils.h"
#include "platform/device.h"
namespace xllm::kernel::cuda {
torch::Tensor cutlass_fused_moe(
const torch::Tensor& input, // [num_tokens, hidden]
const torch::Tensor& token_selected_experts, // [num_tokens, top_k]
const torch::Tensor& token_final_scales, // [num_tokens, top_k]
const torch::Tensor&
fc1_expert_weights, // [num_experts, inter_dim, hidden]
const torch::Tensor&
fc2_expert_weights, // [num_experts, hidden, inter_dim]
torch::ScalarType output_dtype,
const std::vector<torch::Tensor>& quant_scales,
int32_t tp_size,
int32_t tp_rank,
int32_t ep_size,
int32_t ep_rank,
int32_t cluster_size,
int32_t cluster_rank,
const std::optional<torch::Tensor>& fc1_expert_biases,
const std::optional<torch::Tensor>& fc2_expert_biases,
const std::optional<torch::Tensor>& input_sf,
const std::optional<torch::Tensor>& swiglu_alpha,
const std::optional<torch::Tensor>& swiglu_beta,
const std::optional<torch::Tensor>& swiglu_limit,
const std::optional<torch::Tensor>& output,
bool enable_alltoall,
bool use_deepseek_fp8_block_scale,
bool use_w4_group_scaling,
bool use_mxfp8_act_scaling,
bool min_latency_mode,
bool use_packed_weights,
int32_t tune_max_num_tokens,
ActivationType activation_type) {
int64_t num_rows = input.size(0);
int64_t hidden_size = fc2_expert_weights.size(1);
if (min_latency_mode) {
num_rows *= fc2_expert_weights.size(0);
}
std::vector<int64_t> output_shape = {num_rows, hidden_size};
torch::Tensor result_output;
if (output.has_value() && output.value().defined()) {
result_output = output.value();
} else {
torch::TensorOptions options = input.options().dtype(output_dtype);
result_output = torch::empty(output_shape, options);
}
std::string fused_moe_uri = "fused_moe";
if (Device::is_support_sm90a()) {
fused_moe_uri += "_90";
} else if (Device::is_support_sm100a() || Device::is_support_sm100f()) {
fused_moe_uri += "_100";
} else if (Device::is_support_sm120a()) {
fused_moe_uri += "_120";
} else {
LOG(FATAL) << "FusedMoE is only supported on sm90, sm100, sm120.";
}
bind_tvmffi_stream_to_current_torch_stream(input.device());
ffi::Module fused_moe_runner =
get_function(fused_moe_uri, "init")(
to_dl_data_type(input.scalar_type()),
to_dl_data_type(fc1_expert_weights.scalar_type()),
to_dl_data_type(output_dtype),
use_deepseek_fp8_block_scale,
use_w4_group_scaling,
use_mxfp8_act_scaling,
use_packed_weights)
.cast<ffi::Module>();
fused_moe_runner->GetFunction("run_moe").value()(
to_ffi_tensor(result_output),
to_ffi_tensor(input),
to_ffi_tensor(token_selected_experts),
to_ffi_optional_tensor(token_final_scales),
to_ffi_tensor(fc1_expert_weights),
to_ffi_optional_tensor(fc1_expert_biases),
to_ffi_tensor(fc2_expert_weights),
to_ffi_optional_tensor(fc2_expert_biases),
to_ffi_optional_array_tensors(quant_scales),
to_ffi_optional_tensor(input_sf),
to_ffi_optional_tensor(swiglu_alpha),
to_ffi_optional_tensor(swiglu_beta),
to_ffi_optional_tensor(swiglu_limit),
tp_size,
tp_rank,
ep_size,
ep_rank,
cluster_size,
cluster_rank,
enable_alltoall,
min_latency_mode,
/*profile_ids=*/ffi::Optional<ffi::Array<int64_t>>(), // TODO: support
// auto tuning
// profile ids
support_pdl(),
activation_type);
return result_output;
}
} // namespace xllm::kernel::cuda

257
ex_engine/csrc/moe/moeTopKFuncs.cuh Normal file
View File

@@ -0,0 +1,257 @@
/*
* Adapted from
* https://github.com/NVIDIA/TensorRT-LLM/blob/v1.3.0rc2/cpp/tensorrt_llm/kernels/moeTopKFuncs.cuh
* Copyright (c) 2026, The vLLM team.
* SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION. All rights
* reserved. SPDX-License-Identifier: Apache-2.0
*
* Licensed under the Apache License, Version 2.0 (the "License");
* you may not use this file except in compliance with the License.
* You may obtain a copy of the License at
*
* http://www.apache.org/licenses/LICENSE-2.0
*
* Unless required by applicable law or agreed to in writing, software
* distributed under the License is distributed on an "AS IS" BASIS,
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
* See the License for the specific language governing permissions and
* limitations under the License.
*/
#pragma once
#include <cooperative_groups.h>
#include <cooperative_groups/reduce.h>
#include <cub/cub.cuh>
namespace vllm {
namespace moe {
namespace reduce_topk {
namespace cg = cooperative_groups;
static constexpr int kWARP_SIZE = 32;
template <typename T_>
struct TopKRedType {
using T = T_;
static_assert(
std::is_same_v<T, float> || std::is_same_v<T, half> ||
std::is_same_v<T, __nv_bfloat16> || std::is_same_v<T, int>,
"Top K reduction only implemented for int, float, float16 and bfloat16");
using TypeCmp = std::conditional_t<sizeof(T) == 4, uint64_t, uint32_t>;
using IdxT = std::conditional_t<sizeof(T) == 4, int32_t, int16_t>;
static constexpr int kMoveBits = (sizeof(T) == 4) ? 32 : 16;
static constexpr int kMaxIdx = 65535;
TypeCmp compValIdx;
static __host__ __device__ inline TypeCmp makeCmpVal(T val, int32_t idx = 0) {
auto valueBits = cub::Traits<T>::TwiddleIn(
reinterpret_cast<typename cub::Traits<T>::UnsignedBits&>(val));
TypeCmp compactTmp = valueBits;
compactTmp = (compactTmp << kMoveBits) | (0xFFFF & (kMaxIdx - idx));
// Use 65535 minus idx to give higher priority to elements with smaller
// indices.
return compactTmp;
}
static __host__ __device__ void unpack(T& value, int32_t& index,
TypeCmp cmp) {
// Since “65535-idx” is always smaller than 65536 and positive, we can
// directly use it as the lower 16 bits
index = kMaxIdx - static_cast<int32_t>((cmp & 0xFFFF));
auto compactTmp = cmp >> kMoveBits;
auto valueBits = cub::Traits<T>::TwiddleOut(
reinterpret_cast<typename cub::Traits<T>::UnsignedBits&>(compactTmp));
value = reinterpret_cast<T&>(valueBits);
}
__host__ __device__ TopKRedType() = default;
__host__ __device__ TopKRedType(T val, int32_t idx)
: compValIdx(makeCmpVal(val, idx)) {}
__host__ __device__ operator TypeCmp() const noexcept { return compValIdx; }
__device__ inline TypeCmp reduce(
cg::thread_block_tile<kWARP_SIZE> const& warp) {
return cg::reduce(warp, compValIdx, cg::greater<TypeCmp>{});
}
};
////////////////////////////////////////////////////////////////////////////////////////////////////
template <int K_, bool Enable_>
struct TopKIdx {
// by default, empty
};
template <int K_>
struct TopKIdx<K_, true> {
static constexpr int K = K_;
int32_t val[K];
};
////////////////////////////////////////////////////////////////////////////////////////////////////
#define TOPK_SWAP(I, J) \
{ \
auto pairMin = min(topK[I].compValIdx, topK[J].compValIdx); \
auto pairMax = max(topK[I].compValIdx, topK[J].compValIdx); \
topK[I].compValIdx = pairMax; \
topK[J].compValIdx = pairMin; \
}
template <int N, typename RedType>
struct Sort;
template <typename RedType>
struct Sort<1, RedType> {
static __device__ void run(RedType* topK) {}
};
template <typename RedType>
struct Sort<2, RedType> {
static __device__ void run(RedType* topK) { TOPK_SWAP(0, 1); }
};
template <typename RedType>
struct Sort<3, RedType> {
static __device__ void run(RedType* topK) {
TOPK_SWAP(0, 1);
TOPK_SWAP(1, 2);
TOPK_SWAP(0, 1);
}
};
template <typename RedType>
struct Sort<4, RedType> {
static __device__ void run(RedType* topK) {
TOPK_SWAP(0, 2);
TOPK_SWAP(1, 3);
TOPK_SWAP(0, 1);
TOPK_SWAP(2, 3);
TOPK_SWAP(1, 2);
}
};
template <int K, typename Type>
__forceinline__ __device__ void reduceTopK(
cg::thread_block_tile<kWARP_SIZE> const& warp, Type (&out)[K],
int32_t (&outIdx)[K], Type value, int32_t idx, Type const minValue,
int actualK = K) {
static_assert(K > 0, "Top K must have K > 0");
static_assert(K < kWARP_SIZE, "Top K must have K < kWARP_SIZE");
using RedType = TopKRedType<Type>;
RedType topK{value, idx};
typename RedType::TypeCmp packedMax{};
#pragma unroll
for (int kk = 0; kk < actualK; ++kk) {
topK =
kk > 0 && packedMax == topK.compValIdx ? RedType{minValue, idx} : topK;
// get the next largest value
packedMax = topK.reduce(warp);
RedType::unpack(out[kk], outIdx[kk], packedMax);
}
};
template <int K, typename Type, int N, bool IsSorted = false>
__device__ void reduceTopKFunc(cg::thread_block_tile<kWARP_SIZE> const& warp,
Type (&out)[K], int32_t (&outIdx)[K],
Type (&value)[N], int32_t (&idx)[N],
Type minValue, int actualK = K) {
static_assert(K > 0, "Top K must have K > 0");
static_assert(K < kWARP_SIZE, "Top K must have K < kWARP_SIZE");
static_assert(N > 0, "Top K must have N > 0");
static_assert(N < 5,
"Only support candidates number less than or equal to 128");
using RedType = TopKRedType<Type>;
RedType topK[N];
#pragma unroll
for (int nn = 0; nn < N; ++nn) {
topK[nn] = RedType{value[nn], idx[nn]};
}
if constexpr (!IsSorted) {
Sort<N, RedType>::run(topK);
}
typename RedType::TypeCmp packedMax{};
#pragma unroll
for (int kk = 0; kk < actualK; ++kk) {
bool update = kk > 0 && packedMax == topK[0].compValIdx;
#pragma unroll
for (int nn = 0; nn < N; ++nn) {
topK[nn] = update && nn == N - 1 ? RedType{minValue, idx[nn]}
: update ? topK[nn + 1]
: topK[nn];
}
// get the next largest value
packedMax = topK[0].reduce(warp);
RedType::unpack(out[kk], outIdx[kk], packedMax);
}
};
template <int K, typename Type, int N>
__forceinline__ __device__ void reduceTopK(
cg::thread_block_tile<kWARP_SIZE> const& warp, Type (&out)[K],
int32_t (&outIdx)[K], Type (&value)[N], int32_t (&idx)[N],
Type const minValue, int actualK = K) {
static_assert(K > 0, "Top K must have K > 0");
static_assert(K < kWARP_SIZE, "Top K must have K < kWARP_SIZE");
static_assert(N > 0, "Top K must have N > 0");
static_assert(
N <= 16,
"Only support candidates number less than or equal to 16*32=512");
static_assert(N <= 4 || N % 4 == 0,
"Only support candidates number is a multiple of 4*32=128 or "
"less than or equal to 4");
using RedType = TopKRedType<Type>;
if constexpr (N <= 4) {
reduceTopKFunc<K, Type, N>(warp, out, outIdx, value, idx, minValue,
actualK);
} else {
constexpr int numLoops = N / 4;
constexpr int numResults = (numLoops * K - 1) / kWARP_SIZE + 1;
Type topKBufferValue[numResults];
int32_t topKBufferIdx[numResults];
int32_t laneIdx = threadIdx.x % kWARP_SIZE;
for (int ii = 0; ii < numResults; ++ii) {
topKBufferValue[ii] = minValue;
topKBufferIdx[ii] = ii * kWARP_SIZE - 1;
}
for (int loop = 0; loop < numLoops; ++loop) {
int start = loop * 4;
Type topKValue[K];
int32_t topKIdx[K];
Type inValue[4];
int32_t inIdx[4];
for (int i = 0; i < 4; ++i) {
inValue[i] = value[start + i];
inIdx[i] = idx[start + i];
}
reduceTopKFunc<K, Type, 4>(warp, topKValue, topKIdx, inValue, inIdx,
minValue, actualK);
int inOffset = laneIdx % K;
if (laneIdx >= loop * K && laneIdx < (loop + 1) * K) {
topKBufferValue[0] = topKValue[inOffset];
topKBufferIdx[0] = topKIdx[inOffset];
}
if (loop == numLoops - 1 && (laneIdx < (numLoops * K - kWARP_SIZE))) {
topKBufferValue[1] = topKValue[inOffset];
topKBufferIdx[1] = topKIdx[inOffset];
}
}
reduceTopKFunc<K, Type, numResults>(warp, out, outIdx, topKBufferValue,
topKBufferIdx, minValue, actualK);
}
};
#undef TOPK_SWAP
} // namespace reduce_topk
} // namespace moe
} // namespace vllm

833
ex_engine/csrc/moe/moe_align_sum_kernels.cu Normal file
View File

@@ -0,0 +1,833 @@
#include <array>
#include <cub/cub.cuh>
#include <cuda_runtime.h>
#include <torch/csrc/stable/macros.h>
#include <torch/csrc/stable/accelerator.h>
#include <torch/csrc/stable/ops.h>
#include <torch/csrc/stable/tensor.h>
#include <torch/headeronly/core/ScalarType.h>
#include "../../cuda_compat.h"
#include "core/math.hpp"
#include "libtorch_stable/dispatch_utils.h"
#include "libtorch_stable/torch_utils.h"
#define CEILDIV(x, y) (((x) + (y) - 1) / (y))
namespace vllm {
namespace moe {
namespace batched_moe_align_block_size {
// Note num_threads needs to be 1024 for BlockScan Reduction in the kernel.
static constexpr int32_t num_threads = 1024;
static constexpr int32_t num_blocks = 1;
__global__ void batched_moe_align_block_size_kernel(
int32_t const num_batches, int32_t const max_tokens_per_batch,
int32_t const block_size, int32_t const* __restrict__ batch_num_tokens,
int32_t* __restrict__ sorted_ids, int32_t* __restrict__ block_ids,
int32_t* __restrict__ num_tokens_post_pad) {
// TODO(varun): This is a naive implementation. Could be optimized.
size_t const batch_id = threadIdx.x;
size_t const stride = blockDim.x * gridDim.x;
int32_t const num_blocks_per_batch =
CEILDIV(max_tokens_per_batch, block_size);
int32_t const sorted_ids_size =
num_blocks_per_batch * num_batches * block_size;
int32_t const block_ids_size = sorted_ids_size / block_size;
int32_t const SENTINEL =
num_batches * max_tokens_per_batch; // To denote invalid entries.
// Initialize sorted_ids
for (size_t i = threadIdx.x; i < sorted_ids_size; i += stride) {
sorted_ids[i] = SENTINEL;
}
// Initialize expert_ids with -1
for (size_t i = threadIdx.x; i < block_ids_size; i += stride) {
block_ids[i] = -1;
}
int32_t b_num_tokens = 0;
if (batch_id < num_batches) {
b_num_tokens = batch_num_tokens[batch_id];
}
int32_t const ceil_b_num_tokens =
CEILDIV(b_num_tokens, block_size) * block_size;
// Compute prefix sum over token counts per expert
using BlockScan = cub::BlockScan<int32_t, 1024>;
__shared__ typename BlockScan::TempStorage temp_storage;
int cumsum_val;
BlockScan(temp_storage).ExclusiveSum(ceil_b_num_tokens, cumsum_val);
__syncthreads();
bool const is_last_batch = batch_id == (num_batches - 1);
if (is_last_batch) {
*num_tokens_post_pad = cumsum_val + ceil_b_num_tokens;
}
if (batch_id < num_batches) {
int32_t const batch_offset = batch_id * max_tokens_per_batch;
for (size_t i = 0; i < b_num_tokens; ++i) {
sorted_ids[cumsum_val + i] = batch_offset + i;
}
int32_t const block_start = cumsum_val / block_size;
int32_t const num_blocks = ceil_b_num_tokens / block_size;
for (size_t i = 0; i < num_blocks; ++i) {
block_ids[block_start + i] = batch_id;
}
}
}
} // namespace batched_moe_align_block_size
template <typename scalar_t>
__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] = 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 >= num_experts) {
continue;
}
if (has_expert_map) {
expert_id = expert_map[expert_id];
// filter invalid experts
if (expert_id == -1) 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 -1
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>
__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] = 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 (has_expert_map) {
expert_id = expert_map[expert_id];
// filter invalid expert
if (expert_id == -1) 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 -1
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 (has_expert_map) {
expert_id = expert_map[expert_id];
// filter invalid expert
if (expert_id == -1) 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>
__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_align_block_size_kernel(
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 topk_num, bool has_expert_map) {
_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, CEILDIV(max_num_tokens_padded, block_size),
0, -1, topk_num, nullptr, has_expert_map);
}
template <typename scalar_t>
__global__ void 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, bool has_expert_map) {
_count_and_sort_expert_tokens(
topk_ids, sorted_token_ids, cumsum_buffer, expert_map, numel, num_experts,
max_num_tokens_padded, nullptr, 0, topk_num, has_expert_map);
}
template <typename scalar_t, int TOPK>
__global__ void moe_sum_kernel(
scalar_t* __restrict__ out, // [..., d]
const scalar_t* __restrict__ input, // [..., topk, d]
const int d) {
const int64_t token_idx = blockIdx.x;
for (int64_t idx = threadIdx.x; idx < d; idx += blockDim.x) {
scalar_t x = 0.0;
#pragma unroll
for (int k = 0; k < TOPK; ++k) {
x += VLLM_LDG(&input[token_idx * TOPK * d + k * d + idx]);
}
out[token_idx * d + idx] = x;
}
}
template <typename scalar_t, int32_t fill_threads>
__global__ void moe_align_block_size_small_batch_expert_kernel(
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 topk_num,
bool has_expert_map) {
_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,
CEILDIV(max_num_tokens_padded, block_size), -1, 0, topk_num, nullptr,
has_expert_map);
}
template <typename scalar_t>
__global__ void moe_lora_align_block_size_kernel(
scalar_t* __restrict__ topk_ids, int32_t* __restrict__ token_lora_mapping,
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];
// Output buffers are indexed by lora_id (in [0, max_loras)). The grid
// iterates one extra slot to accommodate the "-1" entry that
// active_lora_ids may hold in position 0 for mixed base + LoRA batches;
// guard against any other unexpected lora_id >= max_loras to avoid
// out-of-bounds writes. This mirrors the `lora_id >= max_loras` guard in
// the Triton _fused_moe_lora_kernel.
if (lora_id == -1 || lora_id >= max_loras || adapter_enabled[lora_id] == 0) {
return;
}
// Populate the token_mask based on the token-LoRA mapping
int num_tokens = numel / topk_num;
if (threadIdx.x == 0) {
total_tokens_post_pad[lora_id] = 0;
for (int i = 0; i < num_tokens; i++) {
token_mask[(lora_id * num_tokens) + i] =
(int)token_lora_mapping[i] == lora_id;
}
}
__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, 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 max_loras, int32_t* lora_ids, int32_t* adapter_enabled,
bool has_expert_map) {
int lora_idx = blockIdx.x;
int lora_id = lora_ids[lora_idx];
// Same guard rationale as moe_lora_align_block_size_kernel. Additionally
// skip disabled adapter slots: moe_lora_align_block_size_kernel early-returns
// for them and leaves token_mask[lora_id, :] uninitialized (token_mask is
// allocated with torch::empty), so running the sort loop here would traverse
// garbage mask bits and pollute this slot's rows of sorted_token_ids and
// cumsum_buffer. Downstream consumers already skip disabled slots, so the
// pollution is dormant today, but the check keeps behavior symmetric with
// the other two align kernels and avoids O(numel) wasted work per disabled
// slot. Short-circuit evaluation ensures adapter_enabled is only indexed
// after lora_id is confirmed to be in [0, max_loras).
if (lora_id == -1 || lora_id >= max_loras || 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* token_lora_mapping,
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];
// Same guard rationale as moe_lora_align_block_size_kernel.
if (lora_id == -1 || lora_id >= max_loras || adapter_enabled[lora_id] == 0) {
return;
}
int num_tokens = numel / topk_num;
if (threadIdx.x == 0) {
total_tokens_post_pad[lora_id] = 0;
for (int i = 0; i < num_tokens; i++) {
token_mask[(lora_id * num_tokens) + i] =
(int)token_lora_mapping[i] == lora_id;
}
}
__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, lora_id, topk_num, &token_mask[(lora_id * num_tokens)],
has_expert_map);
}
} // namespace moe
} // namespace vllm
// taken from
// https://github.com/sgl-project/sglang/blob/8b5f83ed3b7d2a49ad5c5cd5aa61c5d502f47dbc
void moe_align_block_size(
torch::stable::Tensor topk_ids, int64_t num_experts, int64_t block_size,
torch::stable::Tensor sorted_token_ids, torch::stable::Tensor experts_ids,
torch::stable::Tensor num_tokens_post_pad,
std::optional<torch::stable::Tensor> maybe_expert_map) {
const cudaStream_t stream =
get_current_cuda_stream(topk_ids.get_device_index());
int64_t padded_num_experts =
((num_experts + WARP_SIZE - 1) / WARP_SIZE) * WARP_SIZE;
int experts_per_warp = WARP_SIZE;
int threads = 1024;
threads = ((threads + WARP_SIZE - 1) / WARP_SIZE) * WARP_SIZE;
// BlockScan uses 1024 threads and assigns one thread per expert.
STD_TORCH_CHECK(padded_num_experts < 1024,
"padded_num_experts must be less than 1024");
bool has_expert_map = maybe_expert_map.has_value();
torch::stable::Tensor expert_map;
if (has_expert_map) {
expert_map = maybe_expert_map.value();
} else {
expert_map = torch::stable::new_empty(topk_ids, {0},
torch::headeronly::ScalarType::Int);
}
VLLM_STABLE_DISPATCH_INTEGRAL_AND_UNSIGNED_TYPES(
topk_ids.scalar_type(), "moe_align_block_size_kernel", [&] {
// calc needed amount of shared mem for `cumsum` tensors
bool small_batch_expert_mode =
(topk_ids.numel() < 1024) && (num_experts <= 64);
if (small_batch_expert_mode) {
const int32_t threads = max((int32_t)num_experts, WARP_SIZE);
const int32_t shared_mem_size =
((threads + 1) * num_experts + (num_experts + 1)) *
sizeof(int32_t);
// threadIdx.x >= fill_threads: counting experts and aligning
// threadIdx.x < fill_threads: filling sorted_token_ids
constexpr int32_t fill_threads = 256;
auto small_batch_expert_kernel =
vllm::moe::moe_align_block_size_small_batch_expert_kernel<
scalar_t, fill_threads>;
small_batch_expert_kernel<<<1, fill_threads + threads,
shared_mem_size, stream>>>(
reinterpret_cast<const scalar_t*>(topk_ids.const_data_ptr()),
reinterpret_cast<int32_t*>(sorted_token_ids.mutable_data_ptr()),
reinterpret_cast<int32_t*>(experts_ids.mutable_data_ptr()),
reinterpret_cast<int32_t*>(
num_tokens_post_pad.mutable_data_ptr()),
reinterpret_cast<int32_t*>(expert_map.mutable_data_ptr()),
num_experts, block_size, topk_ids.numel(),
sorted_token_ids.size(0), topk_ids.size(1), has_expert_map);
} else {
torch::stable::Tensor cumsum_buffer = torch::stable::new_empty(
topk_ids, {num_experts + 1}, torch::headeronly::ScalarType::Int);
auto align_kernel = vllm::moe::moe_align_block_size_kernel<scalar_t>;
size_t num_warps = CEILDIV(padded_num_experts, experts_per_warp);
size_t shared_mem_size =
num_warps * experts_per_warp * sizeof(int32_t);
// launch two threadblocks
// blockIdx.x == 0: counting experts and aligning
// blockIdx.x == 1: filling sorted_token_ids
align_kernel<<<2, threads, shared_mem_size, stream>>>(
reinterpret_cast<const scalar_t*>(topk_ids.const_data_ptr()),
reinterpret_cast<int32_t*>(sorted_token_ids.mutable_data_ptr()),
reinterpret_cast<int32_t*>(experts_ids.mutable_data_ptr()),
reinterpret_cast<int32_t*>(
num_tokens_post_pad.mutable_data_ptr()),
reinterpret_cast<int32_t*>(expert_map.mutable_data_ptr()),
num_experts, padded_num_experts, experts_per_warp, block_size,
topk_ids.numel(),
reinterpret_cast<int32_t*>(cumsum_buffer.mutable_data_ptr()),
sorted_token_ids.size(0), topk_ids.size(1), has_expert_map);
const int block_threads = std::min(256, (int)threads);
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(1, actual_blocks);
auto sort_kernel =
vllm::moe::count_and_sort_expert_tokens_kernel<scalar_t>;
sort_kernel<<<gridDims, block_threads, 0, stream>>>(
reinterpret_cast<const scalar_t*>(topk_ids.const_data_ptr()),
reinterpret_cast<int32_t*>(sorted_token_ids.mutable_data_ptr()),
reinterpret_cast<int32_t*>(cumsum_buffer.mutable_data_ptr()),
reinterpret_cast<int32_t*>(expert_map.mutable_data_ptr()),
topk_ids.numel(), num_experts, sorted_token_ids.size(0),
topk_ids.size(1), has_expert_map);
}
});
}
void batched_moe_align_block_size(int64_t max_tokens_per_batch,
int64_t block_size,
const torch::stable::Tensor& batch_num_tokens,
torch::stable::Tensor sorted_ids,
torch::stable::Tensor batch_ids,
torch::stable::Tensor num_tokens_post_pad) {
namespace batched_kernel = vllm::moe::batched_moe_align_block_size;
const cudaStream_t stream =
get_current_cuda_stream(batch_num_tokens.get_device_index());
int32_t const B = batch_num_tokens.size(0);
int32_t const num_blocks_per_batch =
round_to_next_multiple_of(max_tokens_per_batch, block_size) / block_size;
int32_t const num_blocks = num_blocks_per_batch * B;
int64_t const sorted_ids_size = num_blocks * block_size;
STD_TORCH_CHECK(sorted_ids.size(0) == sorted_ids_size);
STD_TORCH_CHECK(batch_ids.size(0) == sorted_ids_size / block_size);
STD_TORCH_CHECK(num_tokens_post_pad.size(0) == 1);
STD_TORCH_CHECK(B <= batched_kernel::num_threads);
batched_kernel::batched_moe_align_block_size_kernel<<<
batched_kernel::num_blocks, batched_kernel::num_threads, 0, stream>>>(
B, max_tokens_per_batch, block_size,
reinterpret_cast<const int32_t*>(batch_num_tokens.const_data_ptr()),
reinterpret_cast<int32_t*>(sorted_ids.mutable_data_ptr()),
reinterpret_cast<int32_t*>(batch_ids.mutable_data_ptr()),
reinterpret_cast<int32_t*>(num_tokens_post_pad.mutable_data_ptr()));
}
void moe_sum(torch::stable::Tensor& input, // [num_tokens, topk, hidden_size]
torch::stable::Tensor& output) // [num_tokens, hidden_size]
{
const int hidden_size = input.size(-1);
const auto num_tokens = output.numel() / hidden_size;
const int topk = input.size(1);
dim3 grid(num_tokens);
dim3 block(std::min(hidden_size, 1024));
const torch::stable::accelerator::DeviceGuard device_guard(
output.get_device_index());
const cudaStream_t stream =
get_current_cuda_stream(output.get_device_index());
switch (topk) {
case 2:
VLLM_STABLE_DISPATCH_FLOATING_TYPES(
input.scalar_type(), "moe_sum_kernel", [&] {
vllm::moe::moe_sum_kernel<scalar_t, 2><<<grid, block, 0, stream>>>(
reinterpret_cast<scalar_t*>(output.mutable_data_ptr()),
reinterpret_cast<const scalar_t*>(input.const_data_ptr()),
hidden_size);
});
break;
case 3:
VLLM_STABLE_DISPATCH_FLOATING_TYPES(
input.scalar_type(), "moe_sum_kernel", [&] {
vllm::moe::moe_sum_kernel<scalar_t, 3><<<grid, block, 0, stream>>>(
reinterpret_cast<scalar_t*>(output.mutable_data_ptr()),
reinterpret_cast<const scalar_t*>(input.const_data_ptr()),
hidden_size);
});
break;
case 4:
VLLM_STABLE_DISPATCH_FLOATING_TYPES(
input.scalar_type(), "moe_sum_kernel", [&] {
vllm::moe::moe_sum_kernel<scalar_t, 4><<<grid, block, 0, stream>>>(
reinterpret_cast<scalar_t*>(output.mutable_data_ptr()),
reinterpret_cast<const scalar_t*>(input.const_data_ptr()),
hidden_size);
});
break;
default:
torch::stable::sum_out(output, input, std::array<int64_t, 1>{1});
break;
}
}
void moe_lora_align_block_size(
torch::stable::Tensor topk_ids, torch::stable::Tensor token_lora_mapping,
int64_t num_experts, int64_t block_size, int64_t max_loras,
int64_t max_num_tokens_padded, int64_t max_num_m_blocks,
torch::stable::Tensor sorted_token_ids, torch::stable::Tensor expert_ids,
torch::stable::Tensor num_tokens_post_pad,
torch::stable::Tensor adapter_enabled, torch::stable::Tensor lora_ids,
std::optional<torch::stable::Tensor> maybe_expert_map) {
const int topk_num = topk_ids.size(1);
STD_TORCH_CHECK(block_size > 0, "block_size should be greater than 0. ");
int device_max_shared_mem;
int dev = topk_ids.get_device_index();
cudaDeviceGetAttribute(&device_max_shared_mem,
cudaDevAttrMaxSharedMemoryPerBlockOptin, dev);
const cudaStream_t stream = get_current_cuda_stream(dev);
int64_t padded_num_experts =
((num_experts + WARP_SIZE - 1) / WARP_SIZE) * WARP_SIZE;
// BlockScan uses 1024 threads and assigns one thread per expert.
STD_TORCH_CHECK(padded_num_experts < 1024,
"padded_num_experts must be less than 1024");
torch::stable::Tensor token_mask =
torch::stable::new_empty(topk_ids, {max_loras * topk_ids.size(0)},
torch::headeronly::ScalarType::Int);
bool has_expert_map = maybe_expert_map.has_value();
torch::stable::Tensor expert_map;
if (has_expert_map) {
expert_map = maybe_expert_map.value();
} else {
expert_map = torch::stable::new_empty(topk_ids, {0},
torch::headeronly::ScalarType::Int);
}
VLLM_STABLE_DISPATCH_INTEGRAL_TYPES(
topk_ids.scalar_type(), "moe_lora_align_sum_kernel", [&] {
bool small_batch_expert_mode =
(topk_ids.numel() < 1024) && (num_experts <= 64);
if (small_batch_expert_mode) {
const int32_t num_thread = 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) {
STD_TORCH_CHECK(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 =
vllm::moe::moe_lora_align_block_size_small_batch_expert_kernel<
scalar_t, fill_threads>;
STD_CUDA_CHECK(VLLM_DevFuncAttribute_SET_MaxDynamicSharedMemorySize(
(void*)kernel, shared_mem));
// Grid size is (max_loras + 1) because active_lora_ids has length
// max_loras + 1: sorted-unique values of token_lora_mapping, which
// can include -1 (base-model tokens) in addition to up to max_loras
// real LoRA slots. Using max_loras would drop the real LoRA slot
// when -1 is present at position 0 and leave output buffers
// uninitialized, causing illegal memory accesses in downstream
// MoE-LoRA kernels. This mirrors the fix made for the Triton
// _fused_moe_lora_kernel grid in vllm-project/vllm#32277.
kernel<<<max_loras + 1, blockDim, shared_mem, stream>>>(
reinterpret_cast<scalar_t*>(topk_ids.mutable_data_ptr()),
reinterpret_cast<int32_t*>(token_lora_mapping.mutable_data_ptr()),
block_size,
reinterpret_cast<int32_t*>(expert_map.mutable_data_ptr()),
num_experts, max_loras, topk_ids.numel(), max_num_tokens_padded,
max_num_m_blocks,
reinterpret_cast<int32_t*>(sorted_token_ids.mutable_data_ptr()),
reinterpret_cast<int32_t*>(expert_ids.mutable_data_ptr()),
topk_num,
reinterpret_cast<int32_t*>(
num_tokens_post_pad.mutable_data_ptr()),
reinterpret_cast<int32_t*>(adapter_enabled.mutable_data_ptr()),
reinterpret_cast<int32_t*>(lora_ids.mutable_data_ptr()),
reinterpret_cast<int32_t*>(token_mask.mutable_data_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);
// cumsum buffer
torch::stable::Tensor cumsum = torch::stable::new_zeros(
topk_ids, {max_loras * (num_experts + 1)},
torch::headeronly::ScalarType::Int);
auto align_kernel =
vllm::moe::moe_lora_align_block_size_kernel<scalar_t>;
// Launch two threadblocks per LoRA slot, across max_loras + 1 slots
// to cover the extra "-1" (base-model tokens) entry that
// active_lora_ids may contain in addition to up to max_loras real
// LoRA slots. Using max_loras would drop the real LoRA slot when -1
// occupies position 0 and leave the output buffers uninitialized,
// causing illegal memory accesses downstream. Mirrors the grid fix
// applied to _fused_moe_lora_kernel in vllm-project/vllm#32277.
// blockIdx.x % 2 == 0: counting experts and aligning
// blockIdx.x % 2 == 1: filling sorted_token_ids
align_kernel<<<(max_loras + 1) * 2, blockDim, shared_mem_size,
stream>>>(
reinterpret_cast<scalar_t*>(topk_ids.mutable_data_ptr()),
reinterpret_cast<int32_t*>(token_lora_mapping.mutable_data_ptr()),
block_size,
reinterpret_cast<int32_t*>(expert_map.mutable_data_ptr()),
num_experts, max_loras, topk_ids.numel(), max_num_tokens_padded,
max_num_m_blocks,
reinterpret_cast<int32_t*>(sorted_token_ids.mutable_data_ptr()),
reinterpret_cast<int32_t*>(expert_ids.mutable_data_ptr()),
topk_num,
reinterpret_cast<int32_t*>(
num_tokens_post_pad.mutable_data_ptr()),
reinterpret_cast<int32_t*>(adapter_enabled.mutable_data_ptr()),
reinterpret_cast<int32_t*>(cumsum.mutable_data_ptr()), WARP_SIZE,
padded_num_experts,
reinterpret_cast<int32_t*>(lora_ids.mutable_data_ptr()),
reinterpret_cast<int32_t*>(token_mask.mutable_data_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);
// Same rationale as align_kernel above: iterate over max_loras + 1
// slots so the sort kernel processes the real LoRA slot even when
// active_lora_ids has -1 at position 0.
dim3 gridDims(max_loras + 1, actual_blocks);
auto sort_kernel =
vllm::moe::lora_count_and_sort_expert_tokens_kernel<scalar_t>;
sort_kernel<<<gridDims, block_threads, 0, stream>>>(
reinterpret_cast<const scalar_t*>(topk_ids.const_data_ptr()),
reinterpret_cast<int32_t*>(sorted_token_ids.mutable_data_ptr()),
reinterpret_cast<int32_t*>(cumsum.mutable_data_ptr()),
reinterpret_cast<int32_t*>(expert_map.mutable_data_ptr()),
topk_ids.numel(), num_experts, max_num_tokens_padded, topk_num,
reinterpret_cast<int32_t*>(token_mask.mutable_data_ptr()),
max_loras,
reinterpret_cast<int32_t*>(lora_ids.mutable_data_ptr()),
reinterpret_cast<int32_t*>(adapter_enabled.mutable_data_ptr()),
has_expert_map);
}
});
}

56
ex_engine/csrc/moe/moe_fused_topk.cu Normal file
View File

@@ -0,0 +1,56 @@
/* Copyright 2026 The xLLM Authors. All Rights Reserved.
Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at
https://github.com/jd-opensource/xllm/blob/main/LICENSE
Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
==============================================================================*/
#include "kernels/cuda/cuda_ops_api.h"
#include "moe_topk_sigmoid_kernels.cuh"
#include "moe_topk_softmax_kernels.cuh"
namespace xllm::kernel::cuda {
std::tuple<torch::Tensor, torch::Tensor> moe_fused_topk(
torch::Tensor& gating_output,
int64_t topk,
bool renormalize,
const std::optional<torch::Tensor>& correction_bias,
const std::string& scoring_func) {
int64_t num_tokens = gating_output.size(0);
torch::Tensor topk_weights = torch::empty(
{num_tokens, topk},
torch::dtype(torch::kFloat32).device(gating_output.device()));
torch::Tensor topk_ids =
torch::empty({num_tokens, topk},
torch::dtype(torch::kInt32).device(gating_output.device()));
if (scoring_func == "softmax") {
std::optional<torch::Tensor> none_correction_bias = std::nullopt;
topk_softmax(topk_weights,
topk_ids,
gating_output,
renormalize,
/*moe_softcapping=*/0.0,
none_correction_bias);
} else if (scoring_func == "sigmoid") {
topk_sigmoid(
topk_weights, topk_ids, gating_output, renormalize, correction_bias);
} else {
LOG(FATAL) << "Unsupported scoring function for moe topk: " << scoring_func
<< "only softmax and sigmoid are supported";
}
return std::make_tuple(topk_weights, topk_ids);
}
} // namespace xllm::kernel::cuda

87
ex_engine/csrc/moe/moe_ops.h Normal file
View File

@@ -0,0 +1,87 @@
#pragma once
#include <torch/csrc/stable/tensor.h>
#include <optional>
#include <tuple>
void topk_softmax(torch::stable::Tensor& topk_weights,
torch::stable::Tensor& topk_indices,
torch::stable::Tensor& token_expert_indices,
torch::stable::Tensor& gating_output, bool renormalize,
std::optional<torch::stable::Tensor> bias);
void topk_sigmoid(torch::stable::Tensor& topk_weights,
torch::stable::Tensor& topk_indices,
torch::stable::Tensor& token_expert_indices,
torch::stable::Tensor& gating_output, bool renormalize,
std::optional<torch::stable::Tensor> bias);
void topk_softplus_sqrt(
torch::stable::Tensor& topk_weights, torch::stable::Tensor& topk_indices,
torch::stable::Tensor& token_expert_indices,
torch::stable::Tensor& gating_output, bool renormalize,
double routed_scaling_factor,
const std::optional<torch::stable::Tensor>& correction_bias,
const std::optional<torch::stable::Tensor>& input_ids,
const std::optional<torch::stable::Tensor>& tid2eid);
void moe_sum(torch::stable::Tensor& input, torch::stable::Tensor& output);
void moe_align_block_size(
torch::stable::Tensor topk_ids, int64_t num_experts, int64_t block_size,
torch::stable::Tensor sorted_token_ids, torch::stable::Tensor experts_ids,
torch::stable::Tensor num_tokens_post_pad,
std::optional<torch::stable::Tensor> maybe_expert_map);
void batched_moe_align_block_size(
int64_t max_tokens_per_batch, int64_t block_size,
const torch::stable::Tensor& expert_num_tokens,
torch::stable::Tensor sorted_ids, torch::stable::Tensor expert_ids,
torch::stable::Tensor num_tokens_post_pad);
void moe_lora_align_block_size(
torch::stable::Tensor topk_ids, torch::stable::Tensor token_lora_mapping,
int64_t num_experts, int64_t block_size, int64_t max_loras,
int64_t max_num_tokens_padded, int64_t max_num_m_blocks,
torch::stable::Tensor sorted_token_ids, torch::stable::Tensor expert_ids,
torch::stable::Tensor num_tokens_post_pad,
torch::stable::Tensor adapter_enabled, torch::stable::Tensor lora_ids,
std::optional<torch::stable::Tensor> maybe_expert_map);
#ifndef USE_ROCM
torch::stable::Tensor moe_wna16_gemm(
torch::stable::Tensor input, torch::stable::Tensor output,
torch::stable::Tensor b_qweight, torch::stable::Tensor b_scales,
std::optional<torch::stable::Tensor> b_qzeros,
std::optional<torch::stable::Tensor> topk_weights,
torch::stable::Tensor sorted_token_ids, torch::stable::Tensor expert_ids,
torch::stable::Tensor num_tokens_post_pad, int64_t top_k,
int64_t BLOCK_SIZE_M, int64_t BLOCK_SIZE_N, int64_t BLOCK_SIZE_K,
int64_t bit);
std::tuple<torch::stable::Tensor, torch::stable::Tensor> grouped_topk(
const torch::stable::Tensor& scores, int64_t n_group, int64_t topk_group,
int64_t topk, bool renormalize, double routed_scaling_factor,
const torch::stable::Tensor& bias, int64_t scoring_func);
#endif
bool moe_permute_unpermute_supported();
int64_t moe_permute_sort_workspace_size(int64_t num_expanded_rows,
int64_t num_expert);
void shuffle_rows(const torch::stable::Tensor& input_tensor,
const torch::stable::Tensor& dst2src_map,
torch::stable::Tensor& output_tensor);
#ifndef USE_ROCM
// DeepSeek V3 optimized router GEMM kernel for SM90+
// Computes output = mat_a @ mat_b.T where:
// mat_a: [num_tokens, hidden_dim] in bf16
// mat_b: [num_experts, hidden_dim] in bf16
// output: [num_tokens, num_experts] in bf16 or fp32
// Supports num_tokens in [1, 16], num_experts in {256, 384}, hidden_dim = 7168
void dsv3_router_gemm(torch::stable::Tensor& output,
const torch::stable::Tensor& mat_a,
const torch::stable::Tensor& mat_b);
#endif

285
ex_engine/csrc/moe/moe_topk.cuh Normal file
View File

@@ -0,0 +1,285 @@
/*
* Copyright (c) 2025, NVIDIA CORPORATION. All rights reserved.
*
* Licensed under the Apache License, Version 2.0 (the "License");
* you may not use this file except in compliance with the License.
* You may obtain a copy of the License at
*
* http://www.apache.org/licenses/LICENSE-2.0
*
* Unless required by applicable law or agreed to in writing, software
* distributed under the License is distributed on an "AS IS" BASIS,
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
* See the License for the specific language governing permissions and
* limitations under the License.
*/
// refers to
// https://github.com/NVIDIA/TensorRT-LLM/blob/main/cpp/tensorrt_llm/kernels/moeTopKFuncs.cuh
#pragma once
#include <cooperative_groups.h>
#include <cooperative_groups/reduce.h>
#include <cub/cub.cuh>
#include "core/kernels/cuda/arch_condition.h"
namespace xllm::kernel::cuda {
namespace reduce_topk {
namespace cg = cooperative_groups;
static constexpr int kWARP_SIZE = 32;
static constexpr bool kTLLM_GEN_HAS_FAST_REDUX = arch::is_major_v<10>;
template <typename T_>
struct TopKRedType {
using T = T_;
static_assert(
std::is_same_v<T, float> || std::is_same_v<T, half> ||
std::is_same_v<T, __nv_bfloat16> || std::is_same_v<T, int>,
"Top K reduction only implemented for int, float, float16 and bfloat16");
using TypeCmp = std::conditional_t<sizeof(T) == 4, uint64_t, uint32_t>;
using IdxT = std::conditional_t<sizeof(T) == 4, int32_t, int16_t>;
static constexpr int kMoveBits = (sizeof(T) == 4) ? 32 : 16;
static constexpr int kMaxIdx = 65535;
TypeCmp compValIdx;
static __host__ __device__ inline TypeCmp makeCmpVal(T val, int32_t idx = 0) {
auto valueBits = cub::Traits<T>::TwiddleIn(
reinterpret_cast<typename cub::Traits<T>::UnsignedBits&>(val));
TypeCmp compactTmp = valueBits;
compactTmp = (compactTmp << kMoveBits) | (0xFFFF & (kMaxIdx - idx));
// Use 65535 minus idx to give higher priority to elements with smaller
// indices.
return compactTmp;
}
static __host__ __device__ void unpack(T& value,
int32_t& index,
TypeCmp cmp) {
// Since “65535-idx” is always smaller than 65536 and positive, we can
// directly use it as the lower 16 bits
index = kMaxIdx - static_cast<int32_t>((cmp & 0xFFFF));
auto compactTmp = cmp >> kMoveBits;
auto valueBits = cub::Traits<T>::TwiddleOut(
reinterpret_cast<typename cub::Traits<T>::UnsignedBits&>(compactTmp));
value = reinterpret_cast<T&>(valueBits);
}
__host__ __device__ TopKRedType() = default;
__host__ __device__ TopKRedType(T val, int32_t idx)
: compValIdx(makeCmpVal(val, idx)) {}
__host__ __device__ operator TypeCmp() const noexcept { return compValIdx; }
__device__ inline TypeCmp reduce(
cg::thread_block_tile<kWARP_SIZE> const& warp) {
if constexpr (!kTLLM_GEN_HAS_FAST_REDUX || sizeof(TypeCmp) == 8) {
return cg::reduce(warp, compValIdx, cg::greater<TypeCmp>{});
} else {
TypeCmp result;
asm("redux.sync.max.u32 %0, %1, 0xffffffff;\n"
: "=r"(result)
: "r"(compValIdx));
return result;
}
}
};
////////////////////////////////////////////////////////////////////////////////////////////////////
template <int K_, bool Enable_>
struct TopKIdx {
// by default, empty
};
template <int K_>
struct TopKIdx<K_, true> {
static constexpr int K = K_;
int32_t val[K];
};
////////////////////////////////////////////////////////////////////////////////////////////////////
#define TOPK_SWAP(I, J) \
{ \
auto pairMin = min(topK[I].compValIdx, topK[J].compValIdx); \
auto pairMax = max(topK[I].compValIdx, topK[J].compValIdx); \
topK[I].compValIdx = pairMax; \
topK[J].compValIdx = pairMin; \
}
template <int N, typename RedType>
struct Sort;
template <typename RedType>
struct Sort<1, RedType> {
static __device__ void run(RedType* topK) {}
};
template <typename RedType>
struct Sort<2, RedType> {
static __device__ void run(RedType* topK) { TOPK_SWAP(0, 1); }
};
template <typename RedType>
struct Sort<3, RedType> {
static __device__ void run(RedType* topK) {
TOPK_SWAP(0, 1);
TOPK_SWAP(1, 2);
TOPK_SWAP(0, 1);
}
};
template <typename RedType>
struct Sort<4, RedType> {
static __device__ void run(RedType* topK) {
TOPK_SWAP(0, 2);
TOPK_SWAP(1, 3);
TOPK_SWAP(0, 1);
TOPK_SWAP(2, 3);
TOPK_SWAP(1, 2);
}
};
template <int K, typename Type>
__forceinline__ __device__ void reduceTopK(
cg::thread_block_tile<kWARP_SIZE> const& warp,
Type (&out)[K],
int32_t (&outIdx)[K],
Type value,
int32_t idx,
Type const minValue,
int actualK = K) {
static_assert(K > 0, "Top K must have K > 0");
static_assert(K < kWARP_SIZE, "Top K must have K < kWARP_SIZE");
using RedType = TopKRedType<Type>;
RedType topK{value, idx};
typename RedType::TypeCmp packedMax{};
#pragma unroll
for (int kk = 0; kk < actualK; ++kk) //@todo: check if actualK is correct
{
topK =
kk > 0 && packedMax == topK.compValIdx ? RedType{minValue, idx} : topK;
// get the next largest value
packedMax = topK.reduce(warp);
RedType::unpack(out[kk], outIdx[kk], packedMax);
}
};
template <int K, typename Type, int N, bool IsSorted = false>
__device__ void reduceTopKFunc(cg::thread_block_tile<kWARP_SIZE> const& warp,
Type (&out)[K],
int32_t (&outIdx)[K],
Type (&value)[N],
int32_t (&idx)[N],
Type minValue,
int actualK = K) {
static_assert(K > 0, "Top K must have K > 0");
static_assert(K < kWARP_SIZE, "Top K must have K < kWARP_SIZE");
static_assert(N > 0, "Top K must have N > 0");
static_assert(N < 5,
"Only support candidates number less than or equal to 128");
using RedType = TopKRedType<Type>;
RedType topK[N];
#pragma unroll
for (int nn = 0; nn < N; ++nn) {
topK[nn] = RedType{value[nn], idx[nn]};
}
if constexpr (!IsSorted) {
Sort<N, RedType>::run(topK);
}
typename RedType::TypeCmp packedMax{};
#pragma unroll
for (int kk = 0; kk < actualK; ++kk) {
bool update = kk > 0 && packedMax == topK[0].compValIdx;
#pragma unroll
for (int nn = 0; nn < N; ++nn) {
topK[nn] = update && nn == N - 1 ? RedType{minValue, idx[nn]}
: update ? topK[nn + 1]
: topK[nn];
}
// get the next largest value
packedMax = topK[0].reduce(warp);
RedType::unpack(out[kk], outIdx[kk], packedMax);
}
};
template <int K, typename Type, int N>
__forceinline__ __device__ void reduceTopK(
cg::thread_block_tile<kWARP_SIZE> const& warp,
Type (&out)[K],
int32_t (&outIdx)[K],
Type (&value)[N],
int32_t (&idx)[N],
Type const minValue,
int actualK = K) {
static_assert(K > 0, "Top K must have K > 0");
static_assert(K < kWARP_SIZE, "Top K must have K < kWARP_SIZE");
static_assert(N > 0, "Top K must have N > 0");
static_assert(
N <= 16,
"Only support candidates number less than or equal to 16*32=512");
static_assert(N <= 4 || N % 4 == 0,
"Only support candidates number is a multiple of 4*32=128 or "
"less than or equal to 4");
using RedType = TopKRedType<Type>;
if constexpr (N <= 4) {
reduceTopKFunc<K, Type, N>(
warp, out, outIdx, value, idx, minValue, actualK);
} else {
constexpr int numLoops = N / 4;
constexpr int numResults = (numLoops * K - 1) / kWARP_SIZE + 1;
Type topKBufferValue[numResults];
int32_t topKBufferIdx[numResults];
int32_t laneIdx = threadIdx.x % kWARP_SIZE;
// Sentinel index must be in [0, kMaxIdx] to survive makeCmpVal pack/unpack
// (kMaxIdx - idx is stored in 16 bits; -1 would become 0 and unpack to
// 65535). Use kMaxIdx so sentinel slots have smallest compValIdx for
// minValue and lose to any real candidate.
for (int ii = 0; ii < numResults; ++ii) {
topKBufferValue[ii] = minValue;
topKBufferIdx[ii] = RedType::kMaxIdx;
}
for (int loop = 0; loop < numLoops; ++loop) {
int start = loop * 4;
Type topKValue[K];
int32_t topKIdx[K];
Type inValue[4];
int32_t inIdx[4];
for (int i = 0; i < 4; ++i) {
inValue[i] = value[start + i];
inIdx[i] = idx[start + i];
}
reduceTopKFunc<K, Type, 4>(
warp, topKValue, topKIdx, inValue, inIdx, minValue, actualK);
int inOffset = laneIdx % K;
if (laneIdx >= loop * K && laneIdx < (loop + 1) * K) {
topKBufferValue[0] = topKValue[inOffset];
topKBufferIdx[0] = topKIdx[inOffset];
}
if (loop == numLoops - 1 && (laneIdx < (numLoops * K - kWARP_SIZE))) {
topKBufferValue[1] = topKValue[inOffset];
topKBufferIdx[1] = topKIdx[inOffset];
}
}
reduceTopKFunc<K, Type, numResults>(
warp, out, outIdx, topKBufferValue, topKBufferIdx, minValue, actualK);
}
};
#undef TOPK_SWAP
} // namespace reduce_topk
} // namespace xllm::kernel::cuda

602
ex_engine/csrc/moe/moe_topk_sigmoid_kernels.cuh Normal file
View File

@@ -0,0 +1,602 @@
// Adapt from
// https://github.com/vllm-project/vllm/blob/v0.7.3/csrc/moe/topk_softmax_kernels.cu
// which is originally adapted from
// https://github.com/NVIDIA/TensorRT-LLM/blob/v0.7.1/cpp/tensorrt_llm/kernels/mixtureOfExperts/moe_kernels.cu
/* Copyright 2025 SGLang Team. All Rights Reserved.
Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at
http://www.apache.org/licenses/LICENSE-2.0
Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
==============================================================================*/
#include <ATen/cuda/CUDAContext.h>
#include <c10/cuda/CUDAGuard.h>
#include <torch/all.h>
#include <cub/util_type.cuh>
#include <cuda/functional>
#include "kernels/cuda/device_utils.cuh"
namespace {
using namespace xllm::kernel::cuda;
// ====================== Sigmoid things ===============================
// We have our own implementation of sigmoid here so we can support transposing
// the output in the sigmoid kernel when we extend this module to support
// expert-choice routing.
template <typename T, int TPB>
__launch_bounds__(TPB) __global__
void moe_sigmoid(const T* input,
const bool* finished,
float* output,
const int num_cols,
const float* correction_bias) {
const int thread_row_offset = blockIdx.x * num_cols;
// Don't touch finished rows.
if ((finished != nullptr) && finished[blockIdx.x]) {
return;
}
// First pass: Apply transformation, find max, and write transformed values to
// output
for (int ii = threadIdx.x; ii < num_cols; ii += TPB) {
const int idx = thread_row_offset + ii;
float val = convert_to_float<T>(input[idx]);
val = 1.0f / (1.0f + expf(-val));
// Apply correction bias if provided
if (correction_bias != nullptr) {
val = val + correction_bias[ii];
}
output[idx] = val; // Store transformed value
}
}
template <int TPB>
__launch_bounds__(TPB) __global__
void moe_topK(const float* inputs_after_sigmoid,
const bool* finished,
float* output,
int* indices,
const int num_experts,
const int k,
const int start_expert,
const int end_expert,
const bool renormalize,
const float* correction_bias) {
using cub_kvp = cub::KeyValuePair<int, float>;
using BlockReduce = cub::BlockReduce<cub_kvp, TPB>;
__shared__ typename BlockReduce::TempStorage tmpStorage;
cub_kvp thread_kvp;
cub::ArgMax arg_max;
const int block_row = blockIdx.x;
const bool row_is_active = finished ? !finished[block_row] : true;
const int thread_read_offset = blockIdx.x * num_experts;
float row_sum_for_renormalize = 0;
for (int k_idx = 0; k_idx < k; ++k_idx) {
thread_kvp.key = 0;
thread_kvp.value = -1.f; // This is OK because inputs are probabilities
cub_kvp inp_kvp;
for (int expert = threadIdx.x; expert < num_experts; expert += TPB) {
const int idx = thread_read_offset + expert;
inp_kvp.key = expert;
inp_kvp.value = inputs_after_sigmoid[idx];
for (int prior_k = 0; prior_k < k_idx; ++prior_k) {
const int prior_winning_expert = indices[k * block_row + prior_k];
if (prior_winning_expert == expert) {
inp_kvp = thread_kvp;
}
}
thread_kvp = arg_max(inp_kvp, thread_kvp);
}
const cub_kvp result_kvp =
BlockReduce(tmpStorage).Reduce(thread_kvp, arg_max);
if (threadIdx.x == 0) {
// Ignore experts the node isn't responsible for with expert parallelism
const int expert = result_kvp.key;
const bool node_uses_expert =
expert >= start_expert && expert < end_expert;
const bool should_process_row = row_is_active && node_uses_expert;
const int idx = k * block_row + k_idx;
float val = result_kvp.value;
if (correction_bias != nullptr) {
val -= correction_bias[expert];
}
output[idx] = val;
indices[idx] = should_process_row ? (expert - start_expert) : num_experts;
assert(indices[idx] >= 0);
row_sum_for_renormalize += val;
}
__syncthreads();
}
if (renormalize && threadIdx.x == 0) {
float row_sum_for_renormalize_inv = 1.f / row_sum_for_renormalize;
for (int k_idx = 0; k_idx < k; ++k_idx) {
const int idx = k * block_row + k_idx;
output[idx] = output[idx] * row_sum_for_renormalize_inv;
}
}
}
// ====================== TopK sigmoid things ===============================
/*
A Top-K gating sigmoid written to exploit when the number of experts in the
MoE layers are a small power of 2. This allows us to cleanly share the rows
among the threads in a single warp and eliminate communication between warps
(so no need to use shared mem).
It fuses the sigmoid, max and argmax into a single kernel.
Limitations:
1) This implementation is intended for when the number of experts is a small
power of 2. 2) This implementation assumes k is small, but will work for any
k.
*/
template <typename T,
int VPT,
int NUM_EXPERTS,
int WARPS_PER_CTA,
int BYTES_PER_LDG>
__launch_bounds__(WARPS_PER_CTA* WARP_SIZE) __global__
void topk_gating_sigmoid(const T* input,
const bool* finished,
float* output,
const int num_rows,
int* indices,
const int k,
const int start_expert,
const int end_expert,
const bool renormalize,
const float* correction_bias) {
// We begin by enforcing compile time assertions and setting up compile time
// constants.
static_assert(VPT == (VPT & -VPT), "VPT must be power of 2");
static_assert(NUM_EXPERTS == (NUM_EXPERTS & -NUM_EXPERTS),
"NUM_EXPERTS must be power of 2");
static_assert(BYTES_PER_LDG == (BYTES_PER_LDG & -BYTES_PER_LDG),
"BYTES_PER_LDG must be power of 2");
static_assert(BYTES_PER_LDG <= 16, "BYTES_PER_LDG must be leq 16");
// Number of bytes each thread pulls in per load
static constexpr int ELTS_PER_LDG = BYTES_PER_LDG / sizeof(T);
static constexpr int ELTS_PER_ROW = NUM_EXPERTS;
static constexpr int THREADS_PER_ROW = ELTS_PER_ROW / VPT;
static constexpr int LDG_PER_THREAD = VPT / ELTS_PER_LDG;
// Restrictions based on previous section.
static_assert(
VPT % ELTS_PER_LDG == 0,
"The elements per thread must be a multiple of the elements per ldg");
static_assert(WARP_SIZE % THREADS_PER_ROW == 0,
"The threads per row must cleanly divide the threads per warp");
static_assert(THREADS_PER_ROW == (THREADS_PER_ROW & -THREADS_PER_ROW),
"THREADS_PER_ROW must be power of 2");
static_assert(THREADS_PER_ROW <= WARP_SIZE,
"THREADS_PER_ROW can be at most warp size");
// We have NUM_EXPERTS elements per row. We specialize for small #experts
static constexpr int ELTS_PER_WARP = WARP_SIZE * VPT;
static constexpr int ROWS_PER_WARP = ELTS_PER_WARP / ELTS_PER_ROW;
static constexpr int ROWS_PER_CTA = WARPS_PER_CTA * ROWS_PER_WARP;
// Restrictions for previous section.
static_assert(ELTS_PER_WARP % ELTS_PER_ROW == 0,
"The elts per row must cleanly divide the total elt per warp");
// ===================== From this point, we finally start computing run-time
// variables. ========================
// Compute CTA and warp rows. We pack multiple rows into a single warp, and a
// block contains WARPS_PER_CTA warps. This, each block processes a chunk of
// rows. We start by computing the start row for each block.
const int cta_base_row = blockIdx.x * ROWS_PER_CTA;
// Now, using the base row per thread block, we compute the base row per warp.
const int warp_base_row = cta_base_row + threadIdx.y * ROWS_PER_WARP;
// The threads in a warp are split into sub-groups that will work on a row.
// We compute row offset for each thread sub-group
const int thread_row_in_warp = threadIdx.x / THREADS_PER_ROW;
const int thread_row = warp_base_row + thread_row_in_warp;
// Threads with indices out of bounds should early exit here.
if (thread_row >= num_rows) {
return;
}
const bool row_is_active = finished ? !finished[thread_row] : true;
// We finally start setting up the read pointers for each thread. First, each
// thread jumps to the start of the row it will read.
const T* thread_row_ptr = input + thread_row * ELTS_PER_ROW;
// Now, we compute the group each thread belong to in order to determine the
// first column to start loads.
const int thread_group_idx = threadIdx.x % THREADS_PER_ROW;
const int first_elt_read_by_thread = thread_group_idx * ELTS_PER_LDG;
const T* thread_read_ptr = thread_row_ptr + first_elt_read_by_thread;
// Determine the pointer type to use to read in the data depending on the
// BYTES_PER_LDG template param. In theory, this can support all powers of 2
// up to 16. NOTE(woosuk): The original implementation uses CUTLASS aligned
// array here. We defined our own aligned array and use it here to avoid the
// dependency on CUTLASS.
using AccessType = AlignedArray<T, ELTS_PER_LDG>;
// Finally, we pull in the data from global mem
T row_chunk_temp[VPT];
AccessType* row_chunk_vec_ptr =
reinterpret_cast<AccessType*>(&row_chunk_temp);
const AccessType* vec_thread_read_ptr =
reinterpret_cast<const AccessType*>(thread_read_ptr);
#pragma unroll
// Note(Byron): interleaved loads to achieve better memory coalescing
// | thread[0] | thread[1] | thread[2] | thread[3] | thread[0] | thread[1] |
// thread[2] | thread[3] | ...
for (int ii = 0; ii < LDG_PER_THREAD; ++ii) {
row_chunk_vec_ptr[ii] = vec_thread_read_ptr[ii * THREADS_PER_ROW];
}
float row_chunk[VPT];
#pragma unroll
// Note(Byron): upcast logits to float32
for (int ii = 0; ii < VPT; ++ii) {
float val = convert_to_float<T>(row_chunk_temp[ii]);
val = 1.0f / (1.0f + expf(-val));
// Apply correction bias if provided
if (correction_bias != nullptr) {
/*
LDG is interleaved
|thread0 LDG| |thread1 LDG| |thread0 LDG| |thread1 LDG|
|--------- group0 --------| |----------group1 --------|
^ local2
*/
const int group_id = ii / ELTS_PER_LDG;
const int local_id = ii % ELTS_PER_LDG;
const int expert_idx = first_elt_read_by_thread +
group_id * THREADS_PER_ROW * ELTS_PER_LDG +
local_id;
val = val + correction_bias[expert_idx];
}
row_chunk[ii] = val;
}
// Now, row_chunk contains the sigmoid of the row chunk. Now, I want to find
// the topk elements in each row, along with the max index.
int start_col = first_elt_read_by_thread;
static constexpr int COLS_PER_GROUP_LDG = ELTS_PER_LDG * THREADS_PER_ROW;
float row_sum_for_renormalize = 0;
for (int k_idx = 0; k_idx < k; ++k_idx) {
// First, each thread does the local argmax
float max_val = row_chunk[0];
int expert = start_col;
#pragma unroll
for (int ldg = 0, col = start_col; ldg < LDG_PER_THREAD;
++ldg, col += COLS_PER_GROUP_LDG) {
#pragma unroll
for (int ii = 0; ii < ELTS_PER_LDG; ++ii) {
float val = row_chunk[ldg * ELTS_PER_LDG + ii];
// No check on the experts here since columns with the smallest index
// are processed first and only updated if > (not >=)
if (val > max_val) {
max_val = val;
expert = col + ii;
}
}
}
// Now, we perform the argmax reduce. We use the butterfly pattern so threads
// reach consensus about the max. This will be useful for K > 1 so that the
// threads can agree on "who" had the max value. That thread can then blank out
// their max with -inf and the warp can run more iterations...
#pragma unroll
for (int mask = THREADS_PER_ROW / 2; mask > 0; mask /= 2) {
float other_max =
XLLM_SHFL_XOR_SYNC_WIDTH(0xffffffff, max_val, mask, THREADS_PER_ROW);
int other_expert =
XLLM_SHFL_XOR_SYNC_WIDTH(0xffffffff, expert, mask, THREADS_PER_ROW);
// We want lower indices to "win" in every thread so we break ties this
// way
if (other_max > max_val ||
(other_max == max_val && other_expert < expert)) {
max_val = other_max;
expert = other_expert;
}
}
// Write the max for this k iteration to global memory.
if (thread_group_idx == 0) {
// Add a guard to ignore experts not included by this node
const bool node_uses_expert =
expert >= start_expert && expert < end_expert;
const bool should_process_row = row_is_active && node_uses_expert;
// The lead thread from each sub-group will write out the final results to
// global memory. (This will be a single) thread per row of the
// input/output matrices.
const int idx = k * thread_row + k_idx;
if (correction_bias != nullptr) {
max_val -= correction_bias[expert];
}
output[idx] = max_val;
indices[idx] = should_process_row ? (expert - start_expert) : NUM_EXPERTS;
row_sum_for_renormalize += max_val;
}
// Finally, we clear the value in the thread with the current max if there
// is another iteration to run.
if (k_idx + 1 < k) {
const int ldg_group_for_expert = expert / COLS_PER_GROUP_LDG;
const int thread_to_clear_in_group =
(expert / ELTS_PER_LDG) % THREADS_PER_ROW;
// Only the thread in the group which produced the max will reset the
// "winning" value to -inf.
if (thread_group_idx == thread_to_clear_in_group) {
const int offset_for_expert = expert % ELTS_PER_LDG;
// Safe to set to any negative value since row_chunk values must be
// between 0 and 1.
row_chunk[ldg_group_for_expert * ELTS_PER_LDG + offset_for_expert] =
-10000.f;
}
}
}
// Fuse renormalization of topk_weights into this kernel
if (renormalize && thread_group_idx == 0) {
float row_sum_for_renormalize_inv = 1.f / row_sum_for_renormalize;
#pragma unroll
for (int k_idx = 0; k_idx < k; ++k_idx) {
const int idx = k * thread_row + k_idx;
output[idx] = output[idx] * row_sum_for_renormalize_inv;
}
}
}
template <typename T, int EXPERTS, int WARPS_PER_TB>
void topk_gating_sigmoid_launcher_helper(const T* input,
const bool* finished,
float* output,
int* indices,
const int num_rows,
const int k,
const int start_expert,
const int end_expert,
const bool renormalize,
const float* correction_bias,
cudaStream_t stream) {
static constexpr std::size_t MAX_BYTES_PER_LDG = 16;
static constexpr int BYTES_PER_LDG =
MIN(MAX_BYTES_PER_LDG, sizeof(T) * EXPERTS);
using Constants = TopkConstants<T, EXPERTS, BYTES_PER_LDG>;
static constexpr int VPT = Constants::VPT;
static constexpr int ROWS_PER_WARP = Constants::ROWS_PER_WARP;
const int num_warps = (num_rows + ROWS_PER_WARP - 1) / ROWS_PER_WARP;
const int num_blocks = (num_warps + WARPS_PER_TB - 1) / WARPS_PER_TB;
dim3 block_dim(WARP_SIZE, WARPS_PER_TB);
topk_gating_sigmoid<T, VPT, EXPERTS, WARPS_PER_TB, BYTES_PER_LDG>
<<<num_blocks, block_dim, 0, stream>>>(input,
finished,
output,
num_rows,
indices,
k,
start_expert,
end_expert,
renormalize,
correction_bias);
}
#define LAUNCH_SIGMOID(TYPE, NUM_EXPERTS, WARPS_PER_TB) \
topk_gating_sigmoid_launcher_helper<TYPE, NUM_EXPERTS, WARPS_PER_TB>( \
gating_output, \
nullptr, \
topk_weights, \
topk_indices, \
num_tokens, \
topk, \
0, \
num_experts, \
renormalize, \
correction_bias, \
stream);
template <typename T>
void topk_gating_sigmoid_kernel_launcher(const T* gating_output,
float* topk_weights,
int* topk_indices,
float* sigmoid_workspace,
const int num_tokens,
const int num_experts,
const int topk,
const bool renormalize,
const float* correction_bias,
cudaStream_t stream) {
static constexpr int WARPS_PER_TB = 4;
switch (num_experts) {
case 1:
LAUNCH_SIGMOID(T, 1, WARPS_PER_TB);
break;
case 2:
LAUNCH_SIGMOID(T, 2, WARPS_PER_TB);
break;
case 4:
LAUNCH_SIGMOID(T, 4, WARPS_PER_TB);
break;
case 8:
LAUNCH_SIGMOID(T, 8, WARPS_PER_TB);
break;
case 16:
LAUNCH_SIGMOID(T, 16, WARPS_PER_TB);
break;
case 32:
LAUNCH_SIGMOID(T, 32, WARPS_PER_TB);
break;
case 64:
LAUNCH_SIGMOID(T, 64, WARPS_PER_TB);
break;
case 128:
LAUNCH_SIGMOID(T, 128, WARPS_PER_TB);
break;
case 256:
LAUNCH_SIGMOID(T, 256, WARPS_PER_TB);
break;
default: {
TORCH_CHECK(sigmoid_workspace != nullptr,
"sigmoid_workspace must be provided for num_experts that are "
"not a power of 2.");
static constexpr int TPB = 256;
moe_sigmoid<T, TPB><<<num_tokens, TPB, 0, stream>>>(gating_output,
nullptr,
sigmoid_workspace,
num_experts,
correction_bias);
moe_topK<TPB><<<num_tokens, TPB, 0, stream>>>(sigmoid_workspace,
nullptr,
topk_weights,
topk_indices,
num_experts,
topk,
0,
num_experts,
renormalize,
correction_bias);
}
}
}
} // namespace
namespace xllm::kernel::cuda {
void topk_sigmoid(torch::Tensor& topk_weights, // [num_tokens, topk]
torch::Tensor& topk_indices, // [num_tokens, topk]
torch::Tensor& gating_output, // [num_tokens, num_experts]
const bool renormalize,
const std::optional<torch::Tensor>& correction_bias) {
// Check data type
CHECK(gating_output.scalar_type() == at::ScalarType::Float ||
gating_output.scalar_type() == at::ScalarType::Half ||
gating_output.scalar_type() == at::ScalarType::BFloat16)
<< "gating_output must be float32, float16, or bfloat16";
// Check dimensions
CHECK(gating_output.dim() == 2)
<< "gating_output must be 2D tensor [num_tokens, num_experts]";
CHECK(topk_weights.dim() == 2)
<< "topk_weights must be 2D tensor [num_tokens, topk]";
CHECK(topk_indices.dim() == 2)
<< "topk_indices must be 2D tensor [num_tokens, topk]";
// Check shapes
CHECK(gating_output.size(0) == topk_weights.size(0))
<< "First dimension of topk_weights must match num_tokens in "
"gating_output";
CHECK(gating_output.size(0) == topk_indices.size(0))
<< "First dimension of topk_indices must match num_tokens in "
"gating_output";
CHECK(topk_weights.size(-1) == topk_indices.size(-1))
<< "Second dimension of topk_indices must match topk in topk_weights";
CHECK(topk_weights.size(-1) <= gating_output.size(-1))
<< "topk must be less than or equal to num_experts";
const int num_experts = static_cast<int>(gating_output.size(-1));
const int num_tokens = static_cast<int>(gating_output.size(0));
const int topk = static_cast<int>(topk_weights.size(-1));
const bool is_pow_2 =
(num_experts != 0) && ((num_experts & (num_experts - 1)) == 0);
const bool needs_workspace = !is_pow_2 || num_experts > 256;
const int64_t workspace_size = needs_workspace ? num_tokens * num_experts : 0;
const at::cuda::OptionalCUDAGuard device_guard(device_of(gating_output));
const cudaStream_t stream = at::cuda::getCurrentCUDAStream();
torch::Tensor sigmoid_workspace = torch::empty(
{workspace_size}, gating_output.options().dtype(at::ScalarType::Float));
const at::ScalarType dtype = gating_output.scalar_type();
// Validate correction_bias if provided - must always be float32
const float* bias_ptr = nullptr;
if (correction_bias.has_value()) {
const torch::Tensor& bias_tensor = correction_bias.value();
CHECK(bias_tensor.dim() == 1)
<< "correction_bias must be 1D tensor [num_experts]";
CHECK(bias_tensor.size(0) == num_experts)
<< "correction_bias size must match num_experts";
CHECK(bias_tensor.scalar_type() == at::ScalarType::Float)
<< "correction_bias must be float32, got " << bias_tensor.scalar_type();
bias_ptr = bias_tensor.data_ptr<float>();
}
if (dtype == at::ScalarType::Float) {
topk_gating_sigmoid_kernel_launcher<float>(
gating_output.data_ptr<float>(),
topk_weights.data_ptr<float>(),
topk_indices.data_ptr<int>(),
sigmoid_workspace.data_ptr<float>(),
num_tokens,
num_experts,
topk,
renormalize,
bias_ptr,
stream);
} else if (dtype == at::ScalarType::Half) {
topk_gating_sigmoid_kernel_launcher<__half>(
reinterpret_cast<const __half*>(gating_output.data_ptr<at::Half>()),
topk_weights.data_ptr<float>(),
topk_indices.data_ptr<int>(),
sigmoid_workspace.data_ptr<float>(),
num_tokens,
num_experts,
topk,
renormalize,
bias_ptr,
stream);
} else if (dtype == at::ScalarType::BFloat16) {
topk_gating_sigmoid_kernel_launcher<__nv_bfloat16>(
reinterpret_cast<const __nv_bfloat16*>(
gating_output.data_ptr<at::BFloat16>()),
topk_weights.data_ptr<float>(),
topk_indices.data_ptr<int>(),
sigmoid_workspace.data_ptr<float>(),
num_tokens,
num_experts,
topk,
renormalize,
bias_ptr,
stream);
} else {
LOG(FATAL) << "Unsupported gating_output dtype: " << dtype;
}
}
} // namespace xllm::kernel::cuda

40
ex_engine/csrc/moe/moe_topk_softmax_kernels.cuh
View File

@@ -22,11 +22,9 @@ limitations under the License.
#include <torch/all.h>
#include <cub/util_type.cuh>
#if CUDA_VERSION >= 12090
#include <cuda/functional>
#endif
#include "device_utils.cuh"
#include "kernels/cuda/device_utils.cuh"
using cub_kvp = cub::KeyValuePair<int, float>;
@@ -707,7 +705,8 @@ void topk_gating_softmax_kernel_launcher(const T* gating_output,
LAUNCH_SOFTMAX(T, 256, WARPS_PER_TB);
break;
default: {
TORCH_CHECK(softmax_workspace != nullptr, "softmax_workspace must be provided for num_experts that are ");
CHECK(softmax_workspace != nullptr)
<< "softmax_workspace must be provided for num_experts that are "
"not a power of 2.";
static constexpr int TPB = 256;
moe_softmax<T, TPB><<<num_tokens, TPB, 0, stream>>>(gating_output,
@@ -753,23 +752,29 @@ void topk_softmax(torch::Tensor& topk_weights, // [num_tokens, topk]
const double moe_softcapping,
const std::optional<torch::Tensor>& correction_bias) {
// Check data type
TORCH_CHECK(gating_output.scalar_type() == at::ScalarType::Float ||
CHECK(gating_output.scalar_type() == at::ScalarType::Float ||
gating_output.scalar_type() == at::ScalarType::Half ||
gating_output.scalar_type() == at::ScalarType::BFloat16,
"gating_output must be float32, float16, or bfloat16");
gating_output.scalar_type() == at::ScalarType::BFloat16)
<< "gating_output must be float32, float16, or bfloat16";
// Check dimensions
TORCH_CHECK(gating_output.dim() == 2, "gating_output must be 2D tensor [num_tokens, num_experts]");
TORCH_CHECK(topk_weights.dim() == 2, "topk_weights must be 2D tensor [num_tokens, topk]");
TORCH_CHECK(topk_indices.dim() == 2, "topk_indices must be 2D tensor [num_tokens, topk]");
CHECK(gating_output.dim() == 2)
<< "gating_output must be 2D tensor [num_tokens, num_experts]";
CHECK(topk_weights.dim() == 2)
<< "topk_weights must be 2D tensor [num_tokens, topk]";
CHECK(topk_indices.dim() == 2)
<< "topk_indices must be 2D tensor [num_tokens, topk]";
// Check shapes
TORCH_CHECK(gating_output.size(0) == topk_weights.size(0), "First dimension of topk_weights must match num_tokens in ");
CHECK(gating_output.size(0) == topk_weights.size(0))
<< "First dimension of topk_weights must match num_tokens in "
"gating_output"
<< "First dimension of topk_indices must match num_tokens in "
"gating_output";
TORCH_CHECK(topk_weights.size(-1) == topk_indices.size(-1), "Second dimension of topk_indices must match topk in topk_weights topk must be less than or equal to num_experts");
CHECK(topk_weights.size(-1) == topk_indices.size(-1))
<< "Second dimension of topk_indices must match topk in topk_weights"
<< "topk must be less than or equal to num_experts";
const int num_experts = static_cast<int>(gating_output.size(-1));
const int num_tokens = static_cast<int>(gating_output.size(0));
@@ -791,9 +796,12 @@ void topk_softmax(torch::Tensor& topk_weights, // [num_tokens, topk]
const float* bias_ptr = nullptr;
if (correction_bias.has_value()) {
const torch::Tensor& bias_tensor = correction_bias.value();
TORCH_CHECK(bias_tensor.dim() == 1, "correction_bias must be 1D tensor [num_experts]");
TORCH_CHECK(bias_tensor.size(0) == num_experts, "correction_bias size must match num_experts");
TORCH_CHECK(bias_tensor.scalar_type() == at::ScalarType::Float, "correction_bias must be float32, got ");
CHECK(bias_tensor.dim() == 1)
<< "correction_bias must be 1D tensor [num_experts]";
CHECK(bias_tensor.size(0) == num_experts)
<< "correction_bias size must match num_experts";
CHECK(bias_tensor.scalar_type() == at::ScalarType::Float)
<< "correction_bias must be float32, got " << bias_tensor.scalar_type();
bias_ptr = bias_tensor.data_ptr<float>();
}
@@ -841,7 +849,7 @@ void topk_softmax(torch::Tensor& topk_weights, // [num_tokens, topk]
bias_ptr,
stream);
} else {
TORCH_CHECK(false, "Unsupported gating_output dtype");
LOG(FATAL) << "Unsupported gating_output dtype: " << dtype;
}
}
} // namespace xllm::kernel::cuda