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

[perf] xllm_fused_qknorm_rope.so compiled+wired, xllm_cache

Browse Source
This commit is contained in:
root
2026-09-01 06:36:44 +00:00
parent 0d88ac4b62
commit b168ff9e2d
19 changed files with 3749 additions and 880 deletions
Download Patch File Download Diff File Expand all files Collapse all files

4
qwen3_6_scripts/ex_engine/build_xllm_kernels.sh
View File

@@ -146,7 +146,7 @@ echo "============================================================"
echo " 6. xllm_moe.so"
echo "============================================================"
build_so "xllm_moe" \
"ex_engine/xllm_kernels/cuda/moe/moe_fused_topk.cu ex_engine/xllm_kernels/cuda/moe/moe_compute_index.cu ex_engine/xllm_kernels/cuda/moe/moe_combine.cu ex_engine/xllm_kernels/cuda/moe/fused_moe.cpp ex_engine/xllm_kernels/cuda/bindings/xllm_moe_bind.cpp"
"ex_engine/xllm_kernels/cuda/moe/moe_fused_topk.cu ex_engine/xllm_kernels/cuda/moe/moe_compute_index.cu ex_engine/xllm_kernels/cuda/moe/moe_combine.cu ex_engine/xllm_kernels/cuda/bindings/xllm_moe_bind.cpp"
echo ""
echo "============================================================"
@@ -154,4 +154,4 @@ echo " Build complete. Output:"
echo "============================================================"
ls -la "${OUTPUT_DIR}"/*.so 2>/dev/null | tail -30
echo ""
echo "Total .so count: $(ls "${OUTPUT_DIR}"/*.so 2>/dev/null | wc -l)"
echo "Total .so count: $(ls "${OUTPUT_DIR}"/*.so 2>/dev/null | wc -l)"

104
qwen3_6_scripts/ex_engine/csrc/moe/moe_topk.cuh
View File

@@ -21,36 +21,73 @@
#pragma once
#include <cooperative_groups.h>
#if !defined(USE_DCU)
#include <cooperative_groups/reduce.h>
#endif
#if defined(USE_MACA)
#include <cuda_bf16.h>
#endif
#if !defined(USE_DCU)
#include <cub/cub.cuh>
#else
#include <hipcub/hipcub.hpp>
#endif
#include "core/kernels/cuda/arch_condition.h"
#include "arch_condition.h"
#if defined(USE_DCU)
#include <hip/hip_bfloat16.h>
#include <hip/hip_fp16.h>
#endif
#include "device_utils.cuh"
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>;
static constexpr int kWarpSize = 32;
#if !defined(USE_DCU)
static constexpr bool kTllmGenHasFastRedux = arch::is_major_v<10>;
#else
static constexpr bool kTllmGenHasFastRedux = false;
#endif
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>,
std::is_same_v<T, BFloat16Type> || 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>;
#if defined(USE_DCU)
using UnsignedBits = std::conditional_t<sizeof(T) == 4, uint32_t, uint16_t>;
#endif
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) {
#if !defined(USE_DCU)
auto valueBits = cub::Traits<T>::TwiddleIn(
reinterpret_cast<typename cub::Traits<T>::UnsignedBits&>(val));
#else
UnsignedBits valueBits = reinterpret_cast<UnsignedBits&>(val);
constexpr UnsignedBits kSignMask =
static_cast<UnsignedBits>(UnsignedBits{1} << (sizeof(T) * 8 - 1));
if constexpr (std::is_same_v<T, int>) {
valueBits = static_cast<UnsignedBits>(valueBits ^ kSignMask);
} else {
valueBits = (valueBits & kSignMask)
? static_cast<UnsignedBits>(~valueBits)
: static_cast<UnsignedBits>(valueBits ^ kSignMask);
}
#endif
TypeCmp compactTmp = valueBits;
compactTmp = (compactTmp << kMoveBits) | (0xFFFF & (kMaxIdx - idx));
// Use 65535 minus idx to give higher priority to elements with smaller
@@ -61,13 +98,26 @@ struct TopKRedType {
static __host__ __device__ void unpack(T& value,
int32_t& index,
TypeCmp cmp) {
// Since “65535-idx” is always smaller than 65536 and positive, we can
// 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;
#if !defined(USE_DCU)
auto valueBits = cub::Traits<T>::TwiddleOut(
reinterpret_cast<typename cub::Traits<T>::UnsignedBits&>(compactTmp));
#else
UnsignedBits valueBits = static_cast<UnsignedBits>(compactTmp);
constexpr UnsignedBits kSignMask =
static_cast<UnsignedBits>(UnsignedBits{1} << (sizeof(T) * 8 - 1));
if constexpr (std::is_same_v<T, int>) {
valueBits = static_cast<UnsignedBits>(valueBits ^ kSignMask);
} else {
valueBits = (valueBits & kSignMask)
? static_cast<UnsignedBits>(valueBits ^ kSignMask)
: static_cast<UnsignedBits>(~valueBits);
}
#endif
value = reinterpret_cast<T&>(valueBits);
}
@@ -79,8 +129,17 @@ struct TopKRedType {
__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) {
cg::thread_block_tile<kWarpSize> const& warp) {
#if defined(USE_DCU)
TypeCmp result = compValIdx;
#pragma unroll
for (int offset = kWarpSize / 2; offset > 0; offset >>= 1) {
TypeCmp other = warp.shfl_down(result, offset);
result = other > result ? other : result;
}
return warp.shfl(result, 0);
#else
if constexpr (!kTllmGenHasFastRedux || sizeof(TypeCmp) == 8) {
return cg::reduce(warp, compValIdx, cg::greater<TypeCmp>{});
} else {
TypeCmp result;
@@ -89,6 +148,7 @@ struct TopKRedType {
: "r"(compValIdx));
return result;
}
#endif
}
};
@@ -150,7 +210,7 @@ struct Sort<4, RedType> {
template <int K, typename Type>
__forceinline__ __device__ void reduceTopK(
cg::thread_block_tile<kWARP_SIZE> const& warp,
cg::thread_block_tile<kWarpSize> const& warp,
Type (&out)[K],
int32_t (&outIdx)[K],
Type value,
@@ -158,7 +218,7 @@ __forceinline__ __device__ void reduceTopK(
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(K < kWarpSize, "Top K must have K < kWarpSize");
using RedType = TopKRedType<Type>;
RedType topK{value, idx};
typename RedType::TypeCmp packedMax{};
@@ -174,7 +234,7 @@ __forceinline__ __device__ void reduceTopK(
};
template <int K, typename Type, int N, bool IsSorted = false>
__device__ void reduceTopKFunc(cg::thread_block_tile<kWARP_SIZE> const& warp,
__device__ void reduceTopKFunc(cg::thread_block_tile<kWarpSize> const& warp,
Type (&out)[K],
int32_t (&outIdx)[K],
Type (&value)[N],
@@ -182,7 +242,7 @@ __device__ void reduceTopKFunc(cg::thread_block_tile<kWARP_SIZE> const& warp,
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(K < kWarpSize, "Top K must have K < kWarpSize");
static_assert(N > 0, "Top K must have N > 0");
static_assert(N < 5,
"Only support candidates number less than or equal to 128");
@@ -214,7 +274,7 @@ __device__ void reduceTopKFunc(cg::thread_block_tile<kWARP_SIZE> const& warp,
template <int K, typename Type, int N>
__forceinline__ __device__ void reduceTopK(
cg::thread_block_tile<kWARP_SIZE> const& warp,
cg::thread_block_tile<kWarpSize> const& warp,
Type (&out)[K],
int32_t (&outIdx)[K],
Type (&value)[N],
@@ -222,7 +282,7 @@ __forceinline__ __device__ void reduceTopK(
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(K < kWarpSize, "Top K must have K < kWarpSize");
static_assert(N > 0, "Top K must have N > 0");
static_assert(
N <= 16,
@@ -236,22 +296,22 @@ __forceinline__ __device__ void reduceTopK(
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;
constexpr int kNumLoops = N / 4;
constexpr int kNumResults = (kNumLoops * K - 1) / kWarpSize + 1;
Type topKBufferValue[numResults];
int32_t topKBufferIdx[numResults];
int32_t laneIdx = threadIdx.x % kWARP_SIZE;
Type topKBufferValue[kNumResults];
int32_t topKBufferIdx[kNumResults];
int32_t laneIdx = threadIdx.x % kWarpSize;
// 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) {
for (int ii = 0; ii < kNumResults; ++ii) {
topKBufferValue[ii] = minValue;
topKBufferIdx[ii] = RedType::kMaxIdx;
}
for (int loop = 0; loop < numLoops; ++loop) {
for (int loop = 0; loop < kNumLoops; ++loop) {
int start = loop * 4;
Type topKValue[K];
int32_t topKIdx[K];
@@ -268,13 +328,13 @@ __forceinline__ __device__ void reduceTopK(
topKBufferValue[0] = topKValue[inOffset];
topKBufferIdx[0] = topKIdx[inOffset];
}
if (loop == numLoops - 1 && (laneIdx < (numLoops * K - kWARP_SIZE))) {
if (loop == kNumLoops - 1 && (laneIdx < (kNumLoops * K - kWarpSize))) {
topKBufferValue[1] = topKValue[inOffset];
topKBufferIdx[1] = topKIdx[inOffset];
}
}
reduceTopKFunc<K, Type, numResults>(
reduceTopKFunc<K, Type, kNumResults>(
warp, out, outIdx, topKBufferValue, topKBufferIdx, minValue, actualK);
}
};

160
qwen3_6_scripts/ex_engine/csrc/moe/moe_topk_sigmoid_kernels.cuh
View File

@@ -22,14 +22,22 @@ limitations under the License.
#include <torch/all.h>
#include <cub/util_type.cuh>
#include <cuda/functional>
#include "kernels/cuda/device_utils.cuh"
#if !defined(USE_DCU) && !defined(USE_MACA)
#endif
#include "device_utils.cuh"
namespace {
using namespace xllm::kernel::cuda;
#if defined(USE_DCU)
static constexpr unsigned long long kSigmoidFullMask = 0xffffffffffffffffULL;
#else
static constexpr unsigned int kSigmoidFullMask = 0xffffffffU;
#endif
// ====================== 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
@@ -183,29 +191,29 @@ __launch_bounds__(WARPS_PER_CTA* WARP_SIZE) __global__
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;
static constexpr int kEltsPerLdg = BYTES_PER_LDG / sizeof(T);
static constexpr int kEltsPerRow = NUM_EXPERTS;
static constexpr int kThreadsPerRow = kEltsPerRow / VPT;
static constexpr int kLdgPerThread = VPT / kEltsPerLdg;
// Restrictions based on previous section.
static_assert(
VPT % ELTS_PER_LDG == 0,
VPT % kEltsPerLdg == 0,
"The elements per thread must be a multiple of the elements per ldg");
static_assert(WARP_SIZE % THREADS_PER_ROW == 0,
static_assert(WARP_SIZE % kThreadsPerRow == 0,
"The threads per row must cleanly divide the threads per warp");
static_assert(THREADS_PER_ROW == (THREADS_PER_ROW & -THREADS_PER_ROW),
static_assert(kThreadsPerRow == (kThreadsPerRow & -kThreadsPerRow),
"THREADS_PER_ROW must be power of 2");
static_assert(THREADS_PER_ROW <= WARP_SIZE,
static_assert(kThreadsPerRow <= 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;
static constexpr int kEltsPerWarp = WARP_SIZE * VPT;
static constexpr int kRowsPerWarp = kEltsPerWarp / kEltsPerRow;
static constexpr int kRowsPerCta = WARPS_PER_CTA * kRowsPerWarp;
// Restrictions for previous section.
static_assert(ELTS_PER_WARP % ELTS_PER_ROW == 0,
static_assert(kEltsPerWarp % kEltsPerRow == 0,
"The elts per row must cleanly divide the total elt per warp");
// ===================== From this point, we finally start computing run-time
@@ -214,14 +222,14 @@ __launch_bounds__(WARPS_PER_CTA* WARP_SIZE) __global__
// 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;
const int cta_base_row = blockIdx.x * kRowsPerCta;
// 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;
const int warp_base_row = cta_base_row + threadIdx.y * kRowsPerWarp;
// 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_in_warp = threadIdx.x / kThreadsPerRow;
const int thread_row = warp_base_row + thread_row_in_warp;
// Threads with indices out of bounds should early exit here.
@@ -232,12 +240,12 @@ __launch_bounds__(WARPS_PER_CTA* WARP_SIZE) __global__
// 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;
const T* thread_row_ptr = input + thread_row * kEltsPerRow;
// 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 int thread_group_idx = threadIdx.x % kThreadsPerRow;
const int first_elt_read_by_thread = thread_group_idx * kEltsPerLdg;
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
@@ -245,7 +253,7 @@ __launch_bounds__(WARPS_PER_CTA* WARP_SIZE) __global__
// 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>;
using AccessType = AlignedArray<T, kEltsPerLdg>;
// Finally, we pull in the data from global mem
T row_chunk_temp[VPT];
@@ -257,8 +265,8 @@ __launch_bounds__(WARPS_PER_CTA* WARP_SIZE) __global__
// 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];
for (int ii = 0; ii < kLdgPerThread; ++ii) {
row_chunk_vec_ptr[ii] = vec_thread_read_ptr[ii * kThreadsPerRow];
}
float row_chunk[VPT];
@@ -275,11 +283,10 @@ __launch_bounds__(WARPS_PER_CTA* WARP_SIZE) __global__
|--------- group0 --------| |----------group1 --------|
^ local2
*/
const int group_id = ii / ELTS_PER_LDG;
const int local_id = ii % ELTS_PER_LDG;
const int group_id = ii / kEltsPerLdg;
const int local_id = ii % kEltsPerLdg;
const int expert_idx = first_elt_read_by_thread +
group_id * THREADS_PER_ROW * ELTS_PER_LDG +
local_id;
group_id * kThreadsPerRow * kEltsPerLdg + local_id;
val = val + correction_bias[expert_idx];
}
@@ -289,7 +296,7 @@ __launch_bounds__(WARPS_PER_CTA* WARP_SIZE) __global__
// 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;
static constexpr int kColsPerGroupLdg = kEltsPerLdg * kThreadsPerRow;
float row_sum_for_renormalize = 0;
@@ -298,11 +305,11 @@ __launch_bounds__(WARPS_PER_CTA* WARP_SIZE) __global__
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) {
for (int ldg = 0, col = start_col; ldg < kLdgPerThread;
++ldg, col += kColsPerGroupLdg) {
#pragma unroll
for (int ii = 0; ii < ELTS_PER_LDG; ++ii) {
float val = row_chunk[ldg * ELTS_PER_LDG + ii];
for (int ii = 0; ii < kEltsPerLdg; ++ii) {
float val = row_chunk[ldg * kEltsPerLdg + ii];
// No check on the experts here since columns with the smallest index
// are processed first and only updated if > (not >=)
@@ -318,11 +325,11 @@ __launch_bounds__(WARPS_PER_CTA* WARP_SIZE) __global__
// 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);
for (int mask = kThreadsPerRow / 2; mask > 0; mask /= 2) {
float other_max = XLLM_SHFL_XOR_SYNC_WIDTH(
kSigmoidFullMask, max_val, mask, kThreadsPerRow);
int other_expert = XLLM_SHFL_XOR_SYNC_WIDTH(
kSigmoidFullMask, expert, mask, kThreadsPerRow);
// We want lower indices to "win" in every thread so we break ties this
// way
@@ -355,17 +362,17 @@ __launch_bounds__(WARPS_PER_CTA* WARP_SIZE) __global__
// 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 ldg_group_for_expert = expert / kColsPerGroupLdg;
const int thread_to_clear_in_group =
(expert / ELTS_PER_LDG) % THREADS_PER_ROW;
(expert / kEltsPerLdg) % kThreadsPerRow;
// 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;
const int offset_for_expert = expert % kEltsPerLdg;
// 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] =
row_chunk[ldg_group_for_expert * kEltsPerLdg + offset_for_expert] =
-10000.f;
}
}
@@ -394,18 +401,17 @@ void topk_gating_sigmoid_launcher_helper(const T* input,
const bool renormalize,
const float* correction_bias,
cudaStream_t stream) {
static constexpr std::size_t MAX_BYTES_PER_LDG = 16;
static constexpr std::size_t kMaxBytesPerLdg = 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;
static constexpr int kBytesPerLdg = MIN(kMaxBytesPerLdg, sizeof(T) * EXPERTS);
using Constants = TopkConstants<T, EXPERTS, kBytesPerLdg>;
static constexpr int kVpt = Constants::VPT;
static constexpr int kRowsPerWarp = Constants::ROWS_PER_WARP;
const int num_warps = (num_rows + kRowsPerWarp - 1) / kRowsPerWarp;
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>
topk_gating_sigmoid<T, kVpt, EXPERTS, WARPS_PER_TB, kBytesPerLdg>
<<<num_blocks, block_dim, 0, stream>>>(input,
finished,
output,
@@ -443,55 +449,55 @@ void topk_gating_sigmoid_kernel_launcher(const T* gating_output,
const bool renormalize,
const float* correction_bias,
cudaStream_t stream) {
static constexpr int WARPS_PER_TB = 4;
static constexpr int kWarpsPerTb = 4;
switch (num_experts) {
case 1:
LAUNCH_SIGMOID(T, 1, WARPS_PER_TB);
LAUNCH_SIGMOID(T, 1, kWarpsPerTb);
break;
case 2:
LAUNCH_SIGMOID(T, 2, WARPS_PER_TB);
LAUNCH_SIGMOID(T, 2, kWarpsPerTb);
break;
case 4:
LAUNCH_SIGMOID(T, 4, WARPS_PER_TB);
LAUNCH_SIGMOID(T, 4, kWarpsPerTb);
break;
case 8:
LAUNCH_SIGMOID(T, 8, WARPS_PER_TB);
LAUNCH_SIGMOID(T, 8, kWarpsPerTb);
break;
case 16:
LAUNCH_SIGMOID(T, 16, WARPS_PER_TB);
LAUNCH_SIGMOID(T, 16, kWarpsPerTb);
break;
case 32:
LAUNCH_SIGMOID(T, 32, WARPS_PER_TB);
LAUNCH_SIGMOID(T, 32, kWarpsPerTb);
break;
case 64:
LAUNCH_SIGMOID(T, 64, WARPS_PER_TB);
LAUNCH_SIGMOID(T, 64, kWarpsPerTb);
break;
case 128:
LAUNCH_SIGMOID(T, 128, WARPS_PER_TB);
LAUNCH_SIGMOID(T, 128, kWarpsPerTb);
break;
case 256:
LAUNCH_SIGMOID(T, 256, WARPS_PER_TB);
LAUNCH_SIGMOID(T, 256, kWarpsPerTb);
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);
static constexpr int kTpb = 256;
moe_sigmoid<T, kTpb><<<num_tokens, kTpb, 0, stream>>>(gating_output,
nullptr,
sigmoid_workspace,
num_experts,
correction_bias);
moe_topK<kTpb><<<num_tokens, kTpb, 0, stream>>>(sigmoid_workspace,
nullptr,
topk_weights,
topk_indices,
num_experts,
topk,
0,
num_experts,
renormalize,
correction_bias);
}
}
}
@@ -583,8 +589,8 @@ void topk_sigmoid(torch::Tensor& topk_weights, // [num_tokens, topk]
bias_ptr,
stream);
} else if (dtype == at::ScalarType::BFloat16) {
topk_gating_sigmoid_kernel_launcher<__nv_bfloat16>(
reinterpret_cast<const __nv_bfloat16*>(
topk_gating_sigmoid_kernel_launcher<BFloat16Type>(
reinterpret_cast<const BFloat16Type*>(
gating_output.data_ptr<at::BFloat16>()),
topk_weights.data_ptr<float>(),
topk_indices.data_ptr<int>(),

209
qwen3_6_scripts/ex_engine/csrc/moe/moe_topk_softmax_kernels.cuh
View File

@@ -22,9 +22,11 @@ limitations under the License.
#include <torch/all.h>
#include <cub/util_type.cuh>
#include <cuda/functional>
#include "kernels/cuda/device_utils.cuh"
// <cuda/functional> requires CUDA 12+ (libcudacxx); BI-V100 runs CUDA 10.2
// and does not ship that header. The include is unused in this file anyway.
#include "device_utils.cuh"
using cub_kvp = cub::KeyValuePair<int, float>;
@@ -32,6 +34,12 @@ namespace {
using namespace xllm::kernel::cuda;
#if defined(USE_DCU)
static constexpr unsigned long long kSoftmaxFullMask = 0xffffffffffffffffULL;
#else
static constexpr unsigned int kSoftmaxFullMask = 0xffffffffU;
#endif
// ====================== Softmax things ===============================
// We have our own implementation of softmax here so we can support transposing
// the output in the softmax kernel when we extend this module to support
@@ -111,9 +119,10 @@ __launch_bounds__(TPB) __global__
}
namespace moe {
struct TopKPair {
static const int PAIR = 2;
static const int MAX_INDEX = 0;
class TopKPair {
public:
static constexpr int kPair = 2;
static constexpr int kMaxIndex = 0;
cub_kvp max;
cub_kvp secondMax;
@@ -122,7 +131,8 @@ struct TopKPair {
: max(max), secondMax(secondMax) {}
};
struct TopKPairArgMax {
class TopKPairArgMax {
public:
__device__ TopKPairArgMax() {}
__device__ __forceinline__ TopKPair
operator()(const TopKPair& candidate1, const TopKPair& candidate2) const {
@@ -176,8 +186,8 @@ __launch_bounds__(TPB) __global__
const int thread_read_offset = blockIdx.x * num_experts;
float row_sum_for_renormalize = 0;
// Each loop finds the top 2 elements,
// thus requiring only ⌈k/2⌉ loops (calculated as (k + 1) / 2).
for (int k_idx = 0; k_idx < (k + TopKPair::PAIR - 1) / TopKPair::PAIR;
// thus requiring only ceil(k / 2) loops (calculated as (k + 1) / 2).
for (int k_idx = 0; k_idx < (k + TopKPair::kPair - 1) / TopKPair::kPair;
++k_idx) {
// Initializing the top 2 elements by the minimum value.
thread_pair.max.key = 0;
@@ -205,9 +215,11 @@ __launch_bounds__(TPB) __global__
if (threadIdx.x == 0) {
#pragma unroll
// updating 2 elements to the result.
for (int i = 0; i < TopKPair::PAIR; i++) {
if (k_idx * 2 + i >= k) break;
cub_kvp result = (i == TopKPair::MAX_INDEX) ? result_pair.max
for (int i = 0; i < TopKPair::kPair; i++) {
if (k_idx * 2 + i >= k) {
break;
}
cub_kvp result = (i == TopKPair::kMaxIndex) ? result_pair.max
: result_pair.secondMax;
int expert = result.key;
bool node_uses_expert = expert >= start_expert && expert < end_expert;
@@ -343,29 +355,29 @@ __launch_bounds__(WARPS_PER_CTA* WARP_SIZE) __global__
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;
static constexpr int kEltsPerLdg = BYTES_PER_LDG / sizeof(T);
static constexpr int kEltsPerRow = NUM_EXPERTS;
static constexpr int kThreadsPerRow = kEltsPerRow / VPT;
static constexpr int kLdgPerThread = VPT / kEltsPerLdg;
// Restrictions based on previous section.
static_assert(
VPT % ELTS_PER_LDG == 0,
VPT % kEltsPerLdg == 0,
"The elements per thread must be a multiple of the elements per ldg");
static_assert(WARP_SIZE % THREADS_PER_ROW == 0,
static_assert(WARP_SIZE % kThreadsPerRow == 0,
"The threads per row must cleanly divide the threads per warp");
static_assert(THREADS_PER_ROW == (THREADS_PER_ROW & -THREADS_PER_ROW),
static_assert(kThreadsPerRow == (kThreadsPerRow & -kThreadsPerRow),
"THREADS_PER_ROW must be power of 2");
static_assert(THREADS_PER_ROW <= WARP_SIZE,
static_assert(kThreadsPerRow <= 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;
static constexpr int kEltsPerWarp = WARP_SIZE * VPT;
static constexpr int kRowsPerWarp = kEltsPerWarp / kEltsPerRow;
static constexpr int kRowsPerCta = WARPS_PER_CTA * kRowsPerWarp;
// Restrictions for previous section.
static_assert(ELTS_PER_WARP % ELTS_PER_ROW == 0,
static_assert(kEltsPerWarp % kEltsPerRow == 0,
"The elts per row must cleanly divide the total elt per warp");
// ===================== From this point, we finally start computing run-time
@@ -374,14 +386,14 @@ __launch_bounds__(WARPS_PER_CTA* WARP_SIZE) __global__
// 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;
const int cta_base_row = blockIdx.x * kRowsPerCta;
// 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;
const int warp_base_row = cta_base_row + threadIdx.y * kRowsPerWarp;
// 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_in_warp = threadIdx.x / kThreadsPerRow;
const int thread_row = warp_base_row + thread_row_in_warp;
// Threads with indices out of bounds should early exit here.
@@ -392,12 +404,12 @@ __launch_bounds__(WARPS_PER_CTA* WARP_SIZE) __global__
// 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;
const T* thread_row_ptr = input + thread_row * kEltsPerRow;
// 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 int thread_group_idx = threadIdx.x % kThreadsPerRow;
const int first_elt_read_by_thread = thread_group_idx * kEltsPerLdg;
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
@@ -405,7 +417,7 @@ __launch_bounds__(WARPS_PER_CTA* WARP_SIZE) __global__
// 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>;
using AccessType = AlignedArray<T, kEltsPerLdg>;
// Finally, we pull in the data from global mem
T row_chunk_temp[VPT];
@@ -417,8 +429,8 @@ __launch_bounds__(WARPS_PER_CTA* WARP_SIZE) __global__
// 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];
for (int ii = 0; ii < kLdgPerThread; ++ii) {
row_chunk_vec_ptr[ii] = vec_thread_read_ptr[ii * kThreadsPerRow];
}
float row_chunk[VPT];
@@ -447,10 +459,10 @@ __launch_bounds__(WARPS_PER_CTA* WARP_SIZE) __global__
|--------- group0 --------| |----------group1 --------|
^ local2
*/
const int group_id = ii / ELTS_PER_LDG;
const int local_id = ii % ELTS_PER_LDG;
const int group_id = ii / kEltsPerLdg;
const int local_id = ii % kEltsPerLdg;
const int expert_idx = first_elt_read_by_thread +
group_id * THREADS_PER_ROW * ELTS_PER_LDG +
group_id * kThreadsPerRow * kEltsPerLdg +
local_id;
val = val + correction_bias[expert_idx];
}
@@ -475,11 +487,11 @@ __launch_bounds__(WARPS_PER_CTA* WARP_SIZE) __global__
// Now, we find the max within the thread group and distribute among the
// threads. We use a butterfly reduce. lane id: 0-31 within a warp
#pragma unroll
for (int mask = THREADS_PER_ROW / 2; mask > 0; mask /= 2) {
for (int mask = kThreadsPerRow / 2; mask > 0; mask /= 2) {
// butterfly reduce with (lane id ^ mask)
thread_max = max(thread_max,
XLLM_SHFL_XOR_SYNC_WIDTH(
0xffffffff, thread_max, mask, THREADS_PER_ROW));
kSoftmaxFullMask, thread_max, mask, kThreadsPerRow));
}
// From this point, thread max in all the threads have the max within the row.
@@ -495,9 +507,9 @@ __launch_bounds__(WARPS_PER_CTA* WARP_SIZE) __global__
// Now, we perform the sum reduce within each thread group. Similar to the max
// reduce, we use a bufferfly pattern.
#pragma unroll
for (int mask = THREADS_PER_ROW / 2; mask > 0; mask /= 2) {
row_sum +=
XLLM_SHFL_XOR_SYNC_WIDTH(0xffffffff, row_sum, mask, THREADS_PER_ROW);
for (int mask = kThreadsPerRow / 2; mask > 0; mask /= 2) {
row_sum += XLLM_SHFL_XOR_SYNC_WIDTH(
kSoftmaxFullMask, row_sum, mask, kThreadsPerRow);
}
// From this point, all threads have the max and the sum for their rows in the
@@ -520,7 +532,7 @@ __launch_bounds__(WARPS_PER_CTA* WARP_SIZE) __global__
// Now, softmax_res contains the softmax 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;
static constexpr int kColsPerGroupLdg = kEltsPerLdg * kThreadsPerRow;
float row_sum_for_renormalize = 0;
@@ -529,11 +541,11 @@ __launch_bounds__(WARPS_PER_CTA* WARP_SIZE) __global__
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) {
for (int ldg = 0, col = start_col; ldg < kLdgPerThread;
++ldg, col += kColsPerGroupLdg) {
#pragma unroll
for (int ii = 0; ii < ELTS_PER_LDG; ++ii) {
float val = row_chunk[ldg * ELTS_PER_LDG + ii];
for (int ii = 0; ii < kEltsPerLdg; ++ii) {
float val = row_chunk[ldg * kEltsPerLdg + ii];
// No check on the experts here since columns with the smallest index
// are processed first and only updated if > (not >=)
@@ -549,11 +561,11 @@ __launch_bounds__(WARPS_PER_CTA* WARP_SIZE) __global__
// 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);
for (int mask = kThreadsPerRow / 2; mask > 0; mask /= 2) {
float other_max = XLLM_SHFL_XOR_SYNC_WIDTH(
kSoftmaxFullMask, max_val, mask, kThreadsPerRow);
int other_expert = XLLM_SHFL_XOR_SYNC_WIDTH(
kSoftmaxFullMask, expert, mask, kThreadsPerRow);
// We want lower indices to "win" in every thread so we break ties this
// way
@@ -583,17 +595,17 @@ __launch_bounds__(WARPS_PER_CTA* WARP_SIZE) __global__
// 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 ldg_group_for_expert = expert / kColsPerGroupLdg;
const int thread_to_clear_in_group =
(expert / ELTS_PER_LDG) % THREADS_PER_ROW;
(expert / kEltsPerLdg) % kThreadsPerRow;
// 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;
const int offset_for_expert = expert % kEltsPerLdg;
// 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] =
row_chunk[ldg_group_for_expert * kEltsPerLdg + offset_for_expert] =
-10000.f;
}
}
@@ -623,18 +635,17 @@ void topk_gating_softmax_launcher_helper(const T* input,
const float moe_softcapping,
const float* correction_bias,
cudaStream_t stream) {
static constexpr std::size_t MAX_BYTES_PER_LDG = 16;
static constexpr std::size_t kMaxBytesPerLdg = 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;
static constexpr int kBytesPerLdg = MIN(kMaxBytesPerLdg, sizeof(T) * EXPERTS);
using Constants = TopkConstants<T, EXPERTS, kBytesPerLdg>;
static constexpr int kVpt = Constants::VPT;
static constexpr int kRowsPerWarp = Constants::ROWS_PER_WARP;
const int num_warps = (num_rows + kRowsPerWarp - 1) / kRowsPerWarp;
const int num_blocks = (num_warps + WARPS_PER_TB - 1) / WARPS_PER_TB;
dim3 block_dim(WARP_SIZE, WARPS_PER_TB);
topk_gating_softmax<T, VPT, EXPERTS, WARPS_PER_TB, BYTES_PER_LDG>
topk_gating_softmax<T, kVpt, EXPERTS, WARPS_PER_TB, kBytesPerLdg>
<<<num_blocks, block_dim, 0, stream>>>(input,
finished,
output,
@@ -675,69 +686,69 @@ void topk_gating_softmax_kernel_launcher(const T* gating_output,
const float moe_softcapping,
const float* correction_bias,
cudaStream_t stream) {
static constexpr int WARPS_PER_TB = 4;
static constexpr int kWarpsPerTb = 4;
switch (num_experts) {
case 1:
LAUNCH_SOFTMAX(T, 1, WARPS_PER_TB);
LAUNCH_SOFTMAX(T, 1, kWarpsPerTb);
break;
case 2:
LAUNCH_SOFTMAX(T, 2, WARPS_PER_TB);
LAUNCH_SOFTMAX(T, 2, kWarpsPerTb);
break;
case 4:
LAUNCH_SOFTMAX(T, 4, WARPS_PER_TB);
LAUNCH_SOFTMAX(T, 4, kWarpsPerTb);
break;
case 8:
LAUNCH_SOFTMAX(T, 8, WARPS_PER_TB);
LAUNCH_SOFTMAX(T, 8, kWarpsPerTb);
break;
case 16:
LAUNCH_SOFTMAX(T, 16, WARPS_PER_TB);
LAUNCH_SOFTMAX(T, 16, kWarpsPerTb);
break;
case 32:
LAUNCH_SOFTMAX(T, 32, WARPS_PER_TB);
LAUNCH_SOFTMAX(T, 32, kWarpsPerTb);
break;
case 64:
LAUNCH_SOFTMAX(T, 64, WARPS_PER_TB);
LAUNCH_SOFTMAX(T, 64, kWarpsPerTb);
break;
case 128:
LAUNCH_SOFTMAX(T, 128, WARPS_PER_TB);
LAUNCH_SOFTMAX(T, 128, kWarpsPerTb);
break;
case 256:
LAUNCH_SOFTMAX(T, 256, WARPS_PER_TB);
LAUNCH_SOFTMAX(T, 256, kWarpsPerTb);
break;
default: {
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,
nullptr,
softmax_workspace,
num_experts,
moe_softcapping,
correction_bias);
static constexpr int kTpb = 256;
moe_softmax<T, kTpb><<<num_tokens, kTpb, 0, stream>>>(gating_output,
nullptr,
softmax_workspace,
num_experts,
moe_softcapping,
correction_bias);
if (topk == 1) {
// Note: As an optimization for better performance,
// the softmax_workspace is overwritten in-place by both moeTopK and
// moe_topk_fast.
moe_topK<TPB><<<num_tokens, TPB, 0, stream>>>(softmax_workspace,
nullptr,
topk_weights,
topk_indices,
num_experts,
topk,
0,
num_experts,
renormalize);
moe_topK<kTpb><<<num_tokens, kTpb, 0, stream>>>(softmax_workspace,
nullptr,
topk_weights,
topk_indices,
num_experts,
topk,
0,
num_experts,
renormalize);
} else {
moe_topk_fast<TPB><<<num_tokens, TPB, 0, stream>>>(softmax_workspace,
nullptr,
topk_weights,
topk_indices,
num_experts,
topk,
0,
num_experts,
renormalize);
moe_topk_fast<kTpb><<<num_tokens, kTpb, 0, stream>>>(softmax_workspace,
nullptr,
topk_weights,
topk_indices,
num_experts,
topk,
0,
num_experts,
renormalize);
}
}
}
@@ -835,8 +846,8 @@ void topk_softmax(torch::Tensor& topk_weights, // [num_tokens, topk]
bias_ptr,
stream);
} else if (dtype == at::ScalarType::BFloat16) {
topk_gating_softmax_kernel_launcher<__nv_bfloat16>(
reinterpret_cast<const __nv_bfloat16*>(
topk_gating_softmax_kernel_launcher<BFloat16Type>(
reinterpret_cast<const BFloat16Type*>(
gating_output.data_ptr<at::BFloat16>()),
topk_weights.data_ptr<float>(),
topk_indices.data_ptr<int>(),

7
qwen3_6_scripts/ex_engine/kernels/param.h
View File

@@ -817,6 +817,7 @@ struct RandomSampleParams {
};
// Rejection sampling parameters for speculative decoding
// PRD #0 (9e0f1402): Added mode field for explicit greedy/probabilistic dispatch
struct RejectionSampleParams {
// Candidate draft token indices to be verified.
// Shape: [total_draft_tokens]. Dtype: int32.
@@ -849,6 +850,10 @@ struct RejectionSampleParams {
// The maximum number of draft tokens in the batch (max value in
// num_draft_tokens).
int32_t max_spec_len;
// PRD #0 (9e0f1402): Explicit sampling mode — Greedy uses exact-match,
// Probabilistic uses min(1, p/q) acceptance ratio.
// Default: Probabilistic (preserves legacy behavior).
uint8_t draft_sampling_mode = 1; // 0=Greedy, 1=Probabilistic
};
// Masked indexer select paged KV cache parameters
@@ -1438,4 +1443,4 @@ struct ChunkGatedDeltaRuleParams {
// Whether to apply L2 norm to q and k inside the kernel. Default: false.
bool use_qk_l2norm_in_kernel = false;
};
} // namespace xllm::kernel
} // namespace xllm::kernel

9
qwen3_6_scripts/ex_engine/xllm_kernels/cuda/headers/cuda_ops_api.h
View File

@@ -17,7 +17,8 @@ limitations under the License.
#include <ATen/DynamicLibrary.h>
#include <ATen/core/dispatch/Dispatcher.h>
#include <glog/logging.h>
// PRD build fix: replace glog with torch logging (no glog-dev on BI-V100)
#include <c10/util/Logging.h>
#include <optional>
#include <tuple>
@@ -54,7 +55,7 @@ void block_copy(torch::Tensor key_cache_ptrs,
torch::Tensor cum_sum,
int64_t numel_per_block,
torch::ScalarType cache_dtype);
#if !defined(USE_DCU)
#if !defined(USE_DCU) && !defined(__ILUVATAR__)
void batch_prefill(const std::string& uri,
ffi::Array<int64_t> plan_info,
torch::Tensor float_workspace_buffer,
@@ -142,7 +143,7 @@ void batch_decode(const std::string& uri,
std::optional<torch::Tensor>& output_lse,
bool use_tensor_core,
std::optional<torch::Tensor> qo_indptr = std::nullopt);
#endif // !defined(USE_DCU)
#endif // !defined(USE_DCU) && !defined(__ILUVATAR__)
void rms_norm(torch::Tensor output,
torch::Tensor input,
torch::Tensor weight,
@@ -303,4 +304,4 @@ torch::Tensor moe_combine_result(const torch::Tensor& gemm2,
int64_t N,
int32_t topk);
} // namespace xllm::kernel::cuda
} // namespace xllm::kernel::cuda

2
qwen3_6_scripts/ex_engine/xllm_kernels/cuda/headers/type_convert.cuh
View File

@@ -57,7 +57,7 @@ class _typeConvert<float> {
};
#if defined(USE_DCU) || (defined(CUDA_VERSION) && (CUDA_VERSION >= 12000)) || \
defined(USE_MACA)
defined(USE_MACA) || defined(__ILUVATAR__)
// CUDA < 12.0 runs into issues with packed type conversion
template <>
class _typeConvert<c10::Half> {

13
qwen3_6_scripts/ex_engine/xllm_kernels/cuda/headers/utils.h
View File

@@ -21,9 +21,10 @@ limitations under the License.
#else
#include <c10/cuda/CUDAGuard.h>
#endif
#include <glog/logging.h>
// PRD build fix: replace glog with c10 logging (no glog-dev on BI-V100)
#include <c10/util/Logging.h>
#include <torch/torch.h>
#if !defined(USE_DCU)
#if !defined(USE_DCU) && !defined(__ILUVATAR__)
#include <tvm/ffi/container/array.h>
#include <tvm/ffi/container/tensor.h>
#include <tvm/ffi/extra/c_env_api.h>
@@ -46,7 +47,7 @@ limitations under the License.
#define HOST_INLINE inline
#endif
#if !defined(USE_DCU)
#if !defined(USE_DCU) && !defined(__ILUVATAR__)
namespace ffi = tvm::ffi;
#endif
@@ -124,7 +125,7 @@ std::string get_batch_decode_uri(torch::ScalarType dtype_q,
std::tuple<torch::Tensor, double> split_scale_param(const torch::Tensor& scale);
#if !defined(USE_DCU)
#if !defined(USE_DCU) && !defined(__ILUVATAR__)
DLDataType to_dl_data_type(torch::ScalarType scalar_type);
// below are tvm-ffi related functions
@@ -159,5 +160,5 @@ inline void bind_tvmffi_stream_to_current_torch_stream(
<< " dev=" << device.index();
}
}
#endif // !defined(USE_DCU)
} // namespace xllm::kernel::cuda
#endif // !defined(USE_DCU) && !defined(__ILUVATAR__)
} // namespace xllm::kernel::cuda

117
qwen3_6_scripts/ex_engine/xllm_kernels/cuda/moe/fused_moe.cpp
View File

@@ -13,10 +13,10 @@ 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"
#include "platform/platform.h"
#include "cuda_ops_api.h"
#include "utils.h"
namespace xllm::kernel::cuda {
@@ -51,74 +51,73 @@ torch::Tensor cutlass_fused_moe(
bool use_packed_weights,
int32_t tune_max_num_tokens,
ActivationType activation_type) {
int64_t num_rows = input.size(0);
int64_t num_tokens = input.size(0);
int64_t hidden_size = fc2_expert_weights.size(1);
int64_t inter_dim = fc1_expert_weights.size(1);
int64_t top_k = token_selected_experts.size(1);
int64_t num_rows = num_tokens;
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);
result_output = torch::zeros({num_rows, hidden_size},
input.options().dtype(output_dtype));
}
std::string fused_moe_uri = "fused_moe";
if (Platform::is_support_sm90a()) {
fused_moe_uri += "_90";
} else if (Platform::is_support_sm100a() || Platform::is_support_sm100f()) {
fused_moe_uri += "_100";
} else if (Platform::is_support_sm120a()) {
fused_moe_uri += "_120";
} else {
LOG(FATAL) << "FusedMoE is only supported on sm90, sm100, sm120.";
if (Platform::is_support_ivcore10()) {
// BI-V100 path: per-token expert gather + matmul + SiLU-gate + matmul
// This replaces the tvm ffi CUTLASS path with native PyTorch ops.
for (int64_t t = 0; t < num_tokens; ++t) {
auto token = input[t].unsqueeze(0); // [1, hidden]
torch::Tensor accum = torch::zeros({1, hidden_size},
input.options().dtype(output_dtype));
for (int64_t k = 0; k < top_k; ++k) {
int64_t expert_id = token_selected_experts[t][k].item<int64_t>();
float scale = token_final_scales[t][k].item<float>();
// gate_up = token @ fc1[expert].T → [1, inter_dim]
auto gate_up = torch::mm(token, fc1_expert_weights[expert_id].t());
if (fc1_expert_biases.has_value()) {
gate_up = gate_up + fc1_expert_biases.value()[expert_id];
}
// SwiGLU: split into gate and up, apply silu(gate) * up
torch::Tensor act;
if (activation_type == ActivationType::SWIGLU ||
activation_type == ActivationType::SWIGLU_BIAS) {
auto chunks = gate_up.chunk(2, /*dim=*/-1);
act = torch::silu(chunks[0]) * chunks[1];
} else if (activation_type == ActivationType::SILU) {
act = torch::silu(gate_up);
} else if (activation_type == ActivationType::GELU) {
act = torch::gelu(gate_up);
} else {
act = gate_up; // identity
}
// down = act @ fc2[expert].T → [1, hidden]
auto down = torch::mm(act, fc2_expert_weights[expert_id].t());
if (fc2_expert_biases.has_value()) {
down = down + fc2_expert_biases.value()[expert_id];
}
accum += down.to(output_dtype) * scale;
}
result_output[t] = accum.squeeze(0);
}
return result_output;
}
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);
// Original NVIDIA GPU path (sm90/sm100/sm120) via tvm ffi
TORCH_CHECK(false,
"cutlass_fused_moe: no supported platform. "
"BI-V100 should use ivcore10 path above; "
"NVIDIA GPUs require sm90+.");
return result_output;
}
} // namespace xllm::kernel::cuda
} // namespace xllm::kernel::cuda

4
qwen3_6_scripts/ex_engine/xllm_kernels/cuda/moe/moe_topk_softmax_kernels.cuh
View File

@@ -23,8 +23,8 @@ limitations under the License.
#include <cub/util_type.cuh>
#if !defined(USE_DCU) && !defined(USE_MACA)
#endif
// <cuda/functional> requires CUDA 12+ (libcudacxx); BI-V100 runs CUDA 10.2
// and does not ship that header. The include is unused in this file anyway.
#include "device_utils.cuh"

7
qwen3_6_scripts/ex_engine/xllm_kernels/param.h
View File

@@ -817,6 +817,7 @@ struct RandomSampleParams {
};
// Rejection sampling parameters for speculative decoding
// PRD #0 (9e0f1402): Added mode field for explicit greedy/probabilistic dispatch
struct RejectionSampleParams {
// Candidate draft token indices to be verified.
// Shape: [total_draft_tokens]. Dtype: int32.
@@ -849,6 +850,10 @@ struct RejectionSampleParams {
// The maximum number of draft tokens in the batch (max value in
// num_draft_tokens).
int32_t max_spec_len;
// PRD #0 (9e0f1402): Explicit sampling mode — Greedy uses exact-match,
// Probabilistic uses min(1, p/q) acceptance ratio.
// Default: Probabilistic (preserves legacy behavior).
uint8_t draft_sampling_mode = 1; // 0=Greedy, 1=Probabilistic
};
// Masked indexer select paged KV cache parameters
@@ -1438,4 +1443,4 @@ struct ChunkGatedDeltaRuleParams {
// Whether to apply L2 norm to q and k inside the kernel. Default: false.
bool use_qk_l2norm_in_kernel = false;
};
} // namespace xllm::kernel
} // namespace xllm::kernel