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 " 6. xllm_moe.so"
echo "============================================================" echo "============================================================"
build_so "xllm_moe" \ 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 ""
echo "============================================================" echo "============================================================"
@@ -154,4 +154,4 @@ echo " Build complete. Output:"
echo "============================================================" echo "============================================================"
ls -la "${OUTPUT_DIR}"/*.so 2>/dev/null | tail -30 ls -la "${OUTPUT_DIR}"/*.so 2>/dev/null | tail -30
echo "" 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 #pragma once
#include <cooperative_groups.h> #include <cooperative_groups.h>
#if !defined(USE_DCU)
#include <cooperative_groups/reduce.h> #include <cooperative_groups/reduce.h>
#endif
#if defined(USE_MACA)
#include <cuda_bf16.h>
#endif
#if !defined(USE_DCU)
#include <cub/cub.cuh> #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 xllm::kernel::cuda {
namespace reduce_topk { namespace reduce_topk {
namespace cg = cooperative_groups; namespace cg = cooperative_groups;
static constexpr int kWARP_SIZE = 32; static constexpr int kWarpSize = 32;
static constexpr bool kTLLM_GEN_HAS_FAST_REDUX = arch::is_major_v<10>; #if !defined(USE_DCU)
static constexpr bool kTllmGenHasFastRedux = arch::is_major_v<10>;
#else
static constexpr bool kTllmGenHasFastRedux = false;
#endif
template <typename T_> template <typename T_>
struct TopKRedType { struct TopKRedType {
using T = T_; using T = T_;
static_assert( static_assert(
std::is_same_v<T, float> || std::is_same_v<T, half> || 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"); "Top K reduction only implemented for int, float, float16 and bfloat16");
using TypeCmp = std::conditional_t<sizeof(T) == 4, uint64_t, uint32_t>; using TypeCmp = std::conditional_t<sizeof(T) == 4, uint64_t, uint32_t>;
using IdxT = std::conditional_t<sizeof(T) == 4, int32_t, int16_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 kMoveBits = (sizeof(T) == 4) ? 32 : 16;
static constexpr int kMaxIdx = 65535; static constexpr int kMaxIdx = 65535;
TypeCmp compValIdx; TypeCmp compValIdx;
static __host__ __device__ inline TypeCmp makeCmpVal(T val, int32_t idx = 0) { static __host__ __device__ inline TypeCmp makeCmpVal(T val, int32_t idx = 0) {
#if !defined(USE_DCU)
auto valueBits = cub::Traits<T>::TwiddleIn( auto valueBits = cub::Traits<T>::TwiddleIn(
reinterpret_cast<typename cub::Traits<T>::UnsignedBits&>(val)); 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; TypeCmp compactTmp = valueBits;
compactTmp = (compactTmp << kMoveBits) | (0xFFFF & (kMaxIdx - idx)); compactTmp = (compactTmp << kMoveBits) | (0xFFFF & (kMaxIdx - idx));
// Use 65535 minus idx to give higher priority to elements with smaller // 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, static __host__ __device__ void unpack(T& value,
int32_t& index, int32_t& index,
TypeCmp cmp) { 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 // directly use it as the lower 16 bits
index = kMaxIdx - static_cast<int32_t>((cmp & 0xFFFF)); index = kMaxIdx - static_cast<int32_t>((cmp & 0xFFFF));
auto compactTmp = cmp >> kMoveBits; auto compactTmp = cmp >> kMoveBits;
#if !defined(USE_DCU)
auto valueBits = cub::Traits<T>::TwiddleOut( auto valueBits = cub::Traits<T>::TwiddleOut(
reinterpret_cast<typename cub::Traits<T>::UnsignedBits&>(compactTmp)); 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); value = reinterpret_cast<T&>(valueBits);
} }
@@ -79,8 +129,17 @@ struct TopKRedType {
__host__ __device__ operator TypeCmp() const noexcept { return compValIdx; } __host__ __device__ operator TypeCmp() const noexcept { return compValIdx; }
__device__ inline TypeCmp reduce( __device__ inline TypeCmp reduce(
cg::thread_block_tile<kWARP_SIZE> const& warp) { cg::thread_block_tile<kWarpSize> const& warp) {
if constexpr (!kTLLM_GEN_HAS_FAST_REDUX || sizeof(TypeCmp) == 8) { #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>{}); return cg::reduce(warp, compValIdx, cg::greater<TypeCmp>{});
} else { } else {
TypeCmp result; TypeCmp result;
@@ -89,6 +148,7 @@ struct TopKRedType {
: "r"(compValIdx)); : "r"(compValIdx));
return result; return result;
} }
#endif
} }
}; };
@@ -150,7 +210,7 @@ struct Sort<4, RedType> {
template <int K, typename Type> template <int K, typename Type>
__forceinline__ __device__ void reduceTopK( __forceinline__ __device__ void reduceTopK(
cg::thread_block_tile<kWARP_SIZE> const& warp, cg::thread_block_tile<kWarpSize> const& warp,
Type (&out)[K], Type (&out)[K],
int32_t (&outIdx)[K], int32_t (&outIdx)[K],
Type value, Type value,
@@ -158,7 +218,7 @@ __forceinline__ __device__ void reduceTopK(
Type const minValue, Type const minValue,
int actualK = K) { int actualK = K) {
static_assert(K > 0, "Top K must have K > 0"); 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>; using RedType = TopKRedType<Type>;
RedType topK{value, idx}; RedType topK{value, idx};
typename RedType::TypeCmp packedMax{}; typename RedType::TypeCmp packedMax{};
@@ -174,7 +234,7 @@ __forceinline__ __device__ void reduceTopK(
}; };
template <int K, typename Type, int N, bool IsSorted = false> 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], Type (&out)[K],
int32_t (&outIdx)[K], int32_t (&outIdx)[K],
Type (&value)[N], Type (&value)[N],
@@ -182,7 +242,7 @@ __device__ void reduceTopKFunc(cg::thread_block_tile<kWARP_SIZE> const& warp,
Type minValue, Type minValue,
int actualK = K) { int actualK = K) {
static_assert(K > 0, "Top K must have K > 0"); 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 > 0, "Top K must have N > 0");
static_assert(N < 5, static_assert(N < 5,
"Only support candidates number less than or equal to 128"); "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> template <int K, typename Type, int N>
__forceinline__ __device__ void reduceTopK( __forceinline__ __device__ void reduceTopK(
cg::thread_block_tile<kWARP_SIZE> const& warp, cg::thread_block_tile<kWarpSize> const& warp,
Type (&out)[K], Type (&out)[K],
int32_t (&outIdx)[K], int32_t (&outIdx)[K],
Type (&value)[N], Type (&value)[N],
@@ -222,7 +282,7 @@ __forceinline__ __device__ void reduceTopK(
Type const minValue, Type const minValue,
int actualK = K) { int actualK = K) {
static_assert(K > 0, "Top K must have K > 0"); 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 > 0, "Top K must have N > 0");
static_assert( static_assert(
N <= 16, N <= 16,
@@ -236,22 +296,22 @@ __forceinline__ __device__ void reduceTopK(
reduceTopKFunc<K, Type, N>( reduceTopKFunc<K, Type, N>(
warp, out, outIdx, value, idx, minValue, actualK); warp, out, outIdx, value, idx, minValue, actualK);
} else { } else {
constexpr int numLoops = N / 4; constexpr int kNumLoops = N / 4;
constexpr int numResults = (numLoops * K - 1) / kWARP_SIZE + 1; constexpr int kNumResults = (kNumLoops * K - 1) / kWarpSize + 1;
Type topKBufferValue[numResults]; Type topKBufferValue[kNumResults];
int32_t topKBufferIdx[numResults]; int32_t topKBufferIdx[kNumResults];
int32_t laneIdx = threadIdx.x % kWARP_SIZE; int32_t laneIdx = threadIdx.x % kWarpSize;
// Sentinel index must be in [0, kMaxIdx] to survive makeCmpVal pack/unpack // 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 // (kMaxIdx - idx is stored in 16 bits; -1 would become 0 and unpack to
// 65535). Use kMaxIdx so sentinel slots have smallest compValIdx for // 65535). Use kMaxIdx so sentinel slots have smallest compValIdx for
// minValue and lose to any real candidate. // minValue and lose to any real candidate.
for (int ii = 0; ii < numResults; ++ii) { for (int ii = 0; ii < kNumResults; ++ii) {
topKBufferValue[ii] = minValue; topKBufferValue[ii] = minValue;
topKBufferIdx[ii] = RedType::kMaxIdx; topKBufferIdx[ii] = RedType::kMaxIdx;
} }
for (int loop = 0; loop < numLoops; ++loop) { for (int loop = 0; loop < kNumLoops; ++loop) {
int start = loop * 4; int start = loop * 4;
Type topKValue[K]; Type topKValue[K];
int32_t topKIdx[K]; int32_t topKIdx[K];
@@ -268,13 +328,13 @@ __forceinline__ __device__ void reduceTopK(
topKBufferValue[0] = topKValue[inOffset]; topKBufferValue[0] = topKValue[inOffset];
topKBufferIdx[0] = topKIdx[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]; topKBufferValue[1] = topKValue[inOffset];
topKBufferIdx[1] = topKIdx[inOffset]; topKBufferIdx[1] = topKIdx[inOffset];
} }
} }
reduceTopKFunc<K, Type, numResults>( reduceTopKFunc<K, Type, kNumResults>(
warp, out, outIdx, topKBufferValue, topKBufferIdx, minValue, actualK); 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 <torch/all.h>
#include <cub/util_type.cuh> #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 { namespace {
using namespace xllm::kernel::cuda; 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 =============================== // ====================== Sigmoid things ===============================
// We have our own implementation of sigmoid here so we can support transposing // 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 // 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"); static_assert(BYTES_PER_LDG <= 16, "BYTES_PER_LDG must be leq 16");
// Number of bytes each thread pulls in per load // Number of bytes each thread pulls in per load
static constexpr int ELTS_PER_LDG = BYTES_PER_LDG / sizeof(T); static constexpr int kEltsPerLdg = BYTES_PER_LDG / sizeof(T);
static constexpr int ELTS_PER_ROW = NUM_EXPERTS; static constexpr int kEltsPerRow = NUM_EXPERTS;
static constexpr int THREADS_PER_ROW = ELTS_PER_ROW / VPT; static constexpr int kThreadsPerRow = kEltsPerRow / VPT;
static constexpr int LDG_PER_THREAD = VPT / ELTS_PER_LDG; static constexpr int kLdgPerThread = VPT / kEltsPerLdg;
// Restrictions based on previous section. // Restrictions based on previous section.
static_assert( static_assert(
VPT % ELTS_PER_LDG == 0, VPT % kEltsPerLdg == 0,
"The elements per thread must be a multiple of the elements per ldg"); "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"); "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"); "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"); "THREADS_PER_ROW can be at most warp size");
// We have NUM_EXPERTS elements per row. We specialize for small #experts // We have NUM_EXPERTS elements per row. We specialize for small #experts
static constexpr int ELTS_PER_WARP = WARP_SIZE * VPT; static constexpr int kEltsPerWarp = WARP_SIZE * VPT;
static constexpr int ROWS_PER_WARP = ELTS_PER_WARP / ELTS_PER_ROW; static constexpr int kRowsPerWarp = kEltsPerWarp / kEltsPerRow;
static constexpr int ROWS_PER_CTA = WARPS_PER_CTA * ROWS_PER_WARP; static constexpr int kRowsPerCta = WARPS_PER_CTA * kRowsPerWarp;
// Restrictions for previous section. // 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"); "The elts per row must cleanly divide the total elt per warp");
// ===================== From this point, we finally start computing run-time // ===================== 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 // 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 // block contains WARPS_PER_CTA warps. This, each block processes a chunk of
// rows. We start by computing the start row for each block. // 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. // 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. // 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 // 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; const int thread_row = warp_base_row + thread_row_in_warp;
// Threads with indices out of bounds should early exit here. // 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 // We finally start setting up the read pointers for each thread. First, each
// thread jumps to the start of the row it will read. // 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 // Now, we compute the group each thread belong to in order to determine the
// first column to start loads. // first column to start loads.
const int thread_group_idx = threadIdx.x % THREADS_PER_ROW; const int thread_group_idx = threadIdx.x % kThreadsPerRow;
const int first_elt_read_by_thread = thread_group_idx * ELTS_PER_LDG; const int first_elt_read_by_thread = thread_group_idx * kEltsPerLdg;
const T* thread_read_ptr = thread_row_ptr + first_elt_read_by_thread; 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 // 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 // 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 // array here. We defined our own aligned array and use it here to avoid the
// dependency on CUTLASS. // dependency on CUTLASS.
using AccessType = AlignedArray<T, ELTS_PER_LDG>; using AccessType = AlignedArray<T, kEltsPerLdg>;
// Finally, we pull in the data from global mem // Finally, we pull in the data from global mem
T row_chunk_temp[VPT]; 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 // Note(Byron): interleaved loads to achieve better memory coalescing
// | thread[0] | thread[1] | thread[2] | thread[3] | thread[0] | thread[1] | // | thread[0] | thread[1] | thread[2] | thread[3] | thread[0] | thread[1] |
// thread[2] | thread[3] | ... // thread[2] | thread[3] | ...
for (int ii = 0; ii < LDG_PER_THREAD; ++ii) { for (int ii = 0; ii < kLdgPerThread; ++ii) {
row_chunk_vec_ptr[ii] = vec_thread_read_ptr[ii * THREADS_PER_ROW]; row_chunk_vec_ptr[ii] = vec_thread_read_ptr[ii * kThreadsPerRow];
} }
float row_chunk[VPT]; float row_chunk[VPT];
@@ -275,11 +283,10 @@ __launch_bounds__(WARPS_PER_CTA* WARP_SIZE) __global__
|--------- group0 --------| |----------group1 --------| |--------- group0 --------| |----------group1 --------|
^ local2 ^ local2
*/ */
const int group_id = ii / ELTS_PER_LDG; const int group_id = ii / kEltsPerLdg;
const int local_id = ii % ELTS_PER_LDG; const int local_id = ii % kEltsPerLdg;
const int expert_idx = first_elt_read_by_thread + const int expert_idx = first_elt_read_by_thread +
group_id * THREADS_PER_ROW * ELTS_PER_LDG + group_id * kThreadsPerRow * kEltsPerLdg + local_id;
local_id;
val = val + correction_bias[expert_idx]; 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 // 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. // the topk elements in each row, along with the max index.
int start_col = first_elt_read_by_thread; 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; float row_sum_for_renormalize = 0;
@@ -298,11 +305,11 @@ __launch_bounds__(WARPS_PER_CTA* WARP_SIZE) __global__
float max_val = row_chunk[0]; float max_val = row_chunk[0];
int expert = start_col; int expert = start_col;
#pragma unroll #pragma unroll
for (int ldg = 0, col = start_col; ldg < LDG_PER_THREAD; for (int ldg = 0, col = start_col; ldg < kLdgPerThread;
++ldg, col += COLS_PER_GROUP_LDG) { ++ldg, col += kColsPerGroupLdg) {
#pragma unroll #pragma unroll
for (int ii = 0; ii < ELTS_PER_LDG; ++ii) { for (int ii = 0; ii < kEltsPerLdg; ++ii) {
float val = row_chunk[ldg * ELTS_PER_LDG + ii]; float val = row_chunk[ldg * kEltsPerLdg + ii];
// No check on the experts here since columns with the smallest index // No check on the experts here since columns with the smallest index
// are processed first and only updated if > (not >=) // 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 // 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... // their max with -inf and the warp can run more iterations...
#pragma unroll #pragma unroll
for (int mask = THREADS_PER_ROW / 2; mask > 0; mask /= 2) { for (int mask = kThreadsPerRow / 2; mask > 0; mask /= 2) {
float other_max = float other_max = XLLM_SHFL_XOR_SYNC_WIDTH(
XLLM_SHFL_XOR_SYNC_WIDTH(0xffffffff, max_val, mask, THREADS_PER_ROW); kSigmoidFullMask, max_val, mask, kThreadsPerRow);
int other_expert = int other_expert = XLLM_SHFL_XOR_SYNC_WIDTH(
XLLM_SHFL_XOR_SYNC_WIDTH(0xffffffff, expert, mask, THREADS_PER_ROW); kSigmoidFullMask, expert, mask, kThreadsPerRow);
// We want lower indices to "win" in every thread so we break ties this // We want lower indices to "win" in every thread so we break ties this
// way // 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 // Finally, we clear the value in the thread with the current max if there
// is another iteration to run. // is another iteration to run.
if (k_idx + 1 < k) { 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 = 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 // Only the thread in the group which produced the max will reset the
// "winning" value to -inf. // "winning" value to -inf.
if (thread_group_idx == thread_to_clear_in_group) { 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 // Safe to set to any negative value since row_chunk values must be
// between 0 and 1. // 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; -10000.f;
} }
} }
@@ -394,18 +401,17 @@ void topk_gating_sigmoid_launcher_helper(const T* input,
const bool renormalize, const bool renormalize,
const float* correction_bias, const float* correction_bias,
cudaStream_t stream) { 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 = static constexpr int kBytesPerLdg = MIN(kMaxBytesPerLdg, sizeof(T) * EXPERTS);
MIN(MAX_BYTES_PER_LDG, sizeof(T) * EXPERTS); using Constants = TopkConstants<T, EXPERTS, kBytesPerLdg>;
using Constants = TopkConstants<T, EXPERTS, BYTES_PER_LDG>; static constexpr int kVpt = Constants::VPT;
static constexpr int VPT = Constants::VPT; static constexpr int kRowsPerWarp = Constants::ROWS_PER_WARP;
static constexpr int ROWS_PER_WARP = Constants::ROWS_PER_WARP; const int num_warps = (num_rows + kRowsPerWarp - 1) / kRowsPerWarp;
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; const int num_blocks = (num_warps + WARPS_PER_TB - 1) / WARPS_PER_TB;
dim3 block_dim(WARP_SIZE, 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, <<<num_blocks, block_dim, 0, stream>>>(input,
finished, finished,
output, output,
@@ -443,55 +449,55 @@ void topk_gating_sigmoid_kernel_launcher(const T* gating_output,
const bool renormalize, const bool renormalize,
const float* correction_bias, const float* correction_bias,
cudaStream_t stream) { cudaStream_t stream) {
static constexpr int WARPS_PER_TB = 4; static constexpr int kWarpsPerTb = 4;
switch (num_experts) { switch (num_experts) {
case 1: case 1:
LAUNCH_SIGMOID(T, 1, WARPS_PER_TB); LAUNCH_SIGMOID(T, 1, kWarpsPerTb);
break; break;
case 2: case 2:
LAUNCH_SIGMOID(T, 2, WARPS_PER_TB); LAUNCH_SIGMOID(T, 2, kWarpsPerTb);
break; break;
case 4: case 4:
LAUNCH_SIGMOID(T, 4, WARPS_PER_TB); LAUNCH_SIGMOID(T, 4, kWarpsPerTb);
break; break;
case 8: case 8:
LAUNCH_SIGMOID(T, 8, WARPS_PER_TB); LAUNCH_SIGMOID(T, 8, kWarpsPerTb);
break; break;
case 16: case 16:
LAUNCH_SIGMOID(T, 16, WARPS_PER_TB); LAUNCH_SIGMOID(T, 16, kWarpsPerTb);
break; break;
case 32: case 32:
LAUNCH_SIGMOID(T, 32, WARPS_PER_TB); LAUNCH_SIGMOID(T, 32, kWarpsPerTb);
break; break;
case 64: case 64:
LAUNCH_SIGMOID(T, 64, WARPS_PER_TB); LAUNCH_SIGMOID(T, 64, kWarpsPerTb);
break; break;
case 128: case 128:
LAUNCH_SIGMOID(T, 128, WARPS_PER_TB); LAUNCH_SIGMOID(T, 128, kWarpsPerTb);
break; break;
case 256: case 256:
LAUNCH_SIGMOID(T, 256, WARPS_PER_TB); LAUNCH_SIGMOID(T, 256, kWarpsPerTb);
break; break;
default: { default: {
TORCH_CHECK(sigmoid_workspace != nullptr, TORCH_CHECK(sigmoid_workspace != nullptr,
"sigmoid_workspace must be provided for num_experts that are " "sigmoid_workspace must be provided for num_experts that are "
"not a power of 2."); "not a power of 2.");
static constexpr int TPB = 256; static constexpr int kTpb = 256;
moe_sigmoid<T, TPB><<<num_tokens, TPB, 0, stream>>>(gating_output, moe_sigmoid<T, kTpb><<<num_tokens, kTpb, 0, stream>>>(gating_output,
nullptr, nullptr,
sigmoid_workspace, sigmoid_workspace,
num_experts, num_experts,
correction_bias); correction_bias);
moe_topK<TPB><<<num_tokens, TPB, 0, stream>>>(sigmoid_workspace, moe_topK<kTpb><<<num_tokens, kTpb, 0, stream>>>(sigmoid_workspace,
nullptr, nullptr,
topk_weights, topk_weights,
topk_indices, topk_indices,
num_experts, num_experts,
topk, topk,
0, 0,
num_experts, num_experts,
renormalize, renormalize,
correction_bias); correction_bias);
} }
} }
} }
@@ -583,8 +589,8 @@ void topk_sigmoid(torch::Tensor& topk_weights, // [num_tokens, topk]
bias_ptr, bias_ptr,
stream); stream);
} else if (dtype == at::ScalarType::BFloat16) { } else if (dtype == at::ScalarType::BFloat16) {
topk_gating_sigmoid_kernel_launcher<__nv_bfloat16>( topk_gating_sigmoid_kernel_launcher<BFloat16Type>(
reinterpret_cast<const __nv_bfloat16*>( reinterpret_cast<const BFloat16Type*>(
gating_output.data_ptr<at::BFloat16>()), gating_output.data_ptr<at::BFloat16>()),
topk_weights.data_ptr<float>(), topk_weights.data_ptr<float>(),
topk_indices.data_ptr<int>(), 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 <torch/all.h>
#include <cub/util_type.cuh> #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>; using cub_kvp = cub::KeyValuePair<int, float>;
@@ -32,6 +34,12 @@ namespace {
using namespace xllm::kernel::cuda; 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 =============================== // ====================== Softmax things ===============================
// We have our own implementation of softmax here so we can support transposing // 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 // the output in the softmax kernel when we extend this module to support
@@ -111,9 +119,10 @@ __launch_bounds__(TPB) __global__
} }
namespace moe { namespace moe {
struct TopKPair { class TopKPair {
static const int PAIR = 2; public:
static const int MAX_INDEX = 0; static constexpr int kPair = 2;
static constexpr int kMaxIndex = 0;
cub_kvp max; cub_kvp max;
cub_kvp secondMax; cub_kvp secondMax;
@@ -122,7 +131,8 @@ struct TopKPair {
: max(max), secondMax(secondMax) {} : max(max), secondMax(secondMax) {}
}; };
struct TopKPairArgMax { class TopKPairArgMax {
public:
__device__ TopKPairArgMax() {} __device__ TopKPairArgMax() {}
__device__ __forceinline__ TopKPair __device__ __forceinline__ TopKPair
operator()(const TopKPair& candidate1, const TopKPair& candidate2) const { 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; const int thread_read_offset = blockIdx.x * num_experts;
float row_sum_for_renormalize = 0; float row_sum_for_renormalize = 0;
// Each loop finds the top 2 elements, // Each loop finds the top 2 elements,
// thus requiring only ⌈k/2⌉ loops (calculated as (k + 1) / 2). // thus requiring only ceil(k / 2) loops (calculated as (k + 1) / 2).
for (int k_idx = 0; k_idx < (k + TopKPair::PAIR - 1) / TopKPair::PAIR; for (int k_idx = 0; k_idx < (k + TopKPair::kPair - 1) / TopKPair::kPair;
++k_idx) { ++k_idx) {
// Initializing the top 2 elements by the minimum value. // Initializing the top 2 elements by the minimum value.
thread_pair.max.key = 0; thread_pair.max.key = 0;
@@ -205,9 +215,11 @@ __launch_bounds__(TPB) __global__
if (threadIdx.x == 0) { if (threadIdx.x == 0) {
#pragma unroll #pragma unroll
// updating 2 elements to the result. // updating 2 elements to the result.
for (int i = 0; i < TopKPair::PAIR; i++) { for (int i = 0; i < TopKPair::kPair; i++) {
if (k_idx * 2 + i >= k) break; if (k_idx * 2 + i >= k) {
cub_kvp result = (i == TopKPair::MAX_INDEX) ? result_pair.max break;
}
cub_kvp result = (i == TopKPair::kMaxIndex) ? result_pair.max
: result_pair.secondMax; : result_pair.secondMax;
int expert = result.key; int expert = result.key;
bool node_uses_expert = expert >= start_expert && expert < end_expert; 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"); static_assert(BYTES_PER_LDG <= 16, "BYTES_PER_LDG must be leq 16");
// Number of bytes each thread pulls in per load // Number of bytes each thread pulls in per load
static constexpr int ELTS_PER_LDG = BYTES_PER_LDG / sizeof(T); static constexpr int kEltsPerLdg = BYTES_PER_LDG / sizeof(T);
static constexpr int ELTS_PER_ROW = NUM_EXPERTS; static constexpr int kEltsPerRow = NUM_EXPERTS;
static constexpr int THREADS_PER_ROW = ELTS_PER_ROW / VPT; static constexpr int kThreadsPerRow = kEltsPerRow / VPT;
static constexpr int LDG_PER_THREAD = VPT / ELTS_PER_LDG; static constexpr int kLdgPerThread = VPT / kEltsPerLdg;
// Restrictions based on previous section. // Restrictions based on previous section.
static_assert( static_assert(
VPT % ELTS_PER_LDG == 0, VPT % kEltsPerLdg == 0,
"The elements per thread must be a multiple of the elements per ldg"); "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"); "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"); "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"); "THREADS_PER_ROW can be at most warp size");
// We have NUM_EXPERTS elements per row. We specialize for small #experts // We have NUM_EXPERTS elements per row. We specialize for small #experts
static constexpr int ELTS_PER_WARP = WARP_SIZE * VPT; static constexpr int kEltsPerWarp = WARP_SIZE * VPT;
static constexpr int ROWS_PER_WARP = ELTS_PER_WARP / ELTS_PER_ROW; static constexpr int kRowsPerWarp = kEltsPerWarp / kEltsPerRow;
static constexpr int ROWS_PER_CTA = WARPS_PER_CTA * ROWS_PER_WARP; static constexpr int kRowsPerCta = WARPS_PER_CTA * kRowsPerWarp;
// Restrictions for previous section. // 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"); "The elts per row must cleanly divide the total elt per warp");
// ===================== From this point, we finally start computing run-time // ===================== 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 // 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 // block contains WARPS_PER_CTA warps. This, each block processes a chunk of
// rows. We start by computing the start row for each block. // 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. // 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. // 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 // 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; const int thread_row = warp_base_row + thread_row_in_warp;
// Threads with indices out of bounds should early exit here. // 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 // We finally start setting up the read pointers for each thread. First, each
// thread jumps to the start of the row it will read. // 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 // Now, we compute the group each thread belong to in order to determine the
// first column to start loads. // first column to start loads.
const int thread_group_idx = threadIdx.x % THREADS_PER_ROW; const int thread_group_idx = threadIdx.x % kThreadsPerRow;
const int first_elt_read_by_thread = thread_group_idx * ELTS_PER_LDG; const int first_elt_read_by_thread = thread_group_idx * kEltsPerLdg;
const T* thread_read_ptr = thread_row_ptr + first_elt_read_by_thread; 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 // 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 // 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 // array here. We defined our own aligned array and use it here to avoid the
// dependency on CUTLASS. // dependency on CUTLASS.
using AccessType = AlignedArray<T, ELTS_PER_LDG>; using AccessType = AlignedArray<T, kEltsPerLdg>;
// Finally, we pull in the data from global mem // Finally, we pull in the data from global mem
T row_chunk_temp[VPT]; 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 // Note(Byron): interleaved loads to achieve better memory coalescing
// | thread[0] | thread[1] | thread[2] | thread[3] | thread[0] | thread[1] | // | thread[0] | thread[1] | thread[2] | thread[3] | thread[0] | thread[1] |
// thread[2] | thread[3] | ... // thread[2] | thread[3] | ...
for (int ii = 0; ii < LDG_PER_THREAD; ++ii) { for (int ii = 0; ii < kLdgPerThread; ++ii) {
row_chunk_vec_ptr[ii] = vec_thread_read_ptr[ii * THREADS_PER_ROW]; row_chunk_vec_ptr[ii] = vec_thread_read_ptr[ii * kThreadsPerRow];
} }
float row_chunk[VPT]; float row_chunk[VPT];
@@ -447,10 +459,10 @@ __launch_bounds__(WARPS_PER_CTA* WARP_SIZE) __global__
|--------- group0 --------| |----------group1 --------| |--------- group0 --------| |----------group1 --------|
^ local2 ^ local2
*/ */
const int group_id = ii / ELTS_PER_LDG; const int group_id = ii / kEltsPerLdg;
const int local_id = ii % ELTS_PER_LDG; const int local_id = ii % kEltsPerLdg;
const int expert_idx = first_elt_read_by_thread + const int expert_idx = first_elt_read_by_thread +
group_id * THREADS_PER_ROW * ELTS_PER_LDG + group_id * kThreadsPerRow * kEltsPerLdg +
local_id; local_id;
val = val + correction_bias[expert_idx]; 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 // 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 // threads. We use a butterfly reduce. lane id: 0-31 within a warp
#pragma unroll #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) // butterfly reduce with (lane id ^ mask)
thread_max = max(thread_max, thread_max = max(thread_max,
XLLM_SHFL_XOR_SYNC_WIDTH( 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. // 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 // Now, we perform the sum reduce within each thread group. Similar to the max
// reduce, we use a bufferfly pattern. // reduce, we use a bufferfly pattern.
#pragma unroll #pragma unroll
for (int mask = THREADS_PER_ROW / 2; mask > 0; mask /= 2) { for (int mask = kThreadsPerRow / 2; mask > 0; mask /= 2) {
row_sum += row_sum += XLLM_SHFL_XOR_SYNC_WIDTH(
XLLM_SHFL_XOR_SYNC_WIDTH(0xffffffff, row_sum, mask, THREADS_PER_ROW); kSoftmaxFullMask, row_sum, mask, kThreadsPerRow);
} }
// From this point, all threads have the max and the sum for their rows in the // 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 // 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. // the topk elements in each row, along with the max index.
int start_col = first_elt_read_by_thread; 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; float row_sum_for_renormalize = 0;
@@ -529,11 +541,11 @@ __launch_bounds__(WARPS_PER_CTA* WARP_SIZE) __global__
float max_val = row_chunk[0]; float max_val = row_chunk[0];
int expert = start_col; int expert = start_col;
#pragma unroll #pragma unroll
for (int ldg = 0, col = start_col; ldg < LDG_PER_THREAD; for (int ldg = 0, col = start_col; ldg < kLdgPerThread;
++ldg, col += COLS_PER_GROUP_LDG) { ++ldg, col += kColsPerGroupLdg) {
#pragma unroll #pragma unroll
for (int ii = 0; ii < ELTS_PER_LDG; ++ii) { for (int ii = 0; ii < kEltsPerLdg; ++ii) {
float val = row_chunk[ldg * ELTS_PER_LDG + ii]; float val = row_chunk[ldg * kEltsPerLdg + ii];
// No check on the experts here since columns with the smallest index // No check on the experts here since columns with the smallest index
// are processed first and only updated if > (not >=) // 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 // 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... // their max with -inf and the warp can run more iterations...
#pragma unroll #pragma unroll
for (int mask = THREADS_PER_ROW / 2; mask > 0; mask /= 2) { for (int mask = kThreadsPerRow / 2; mask > 0; mask /= 2) {
float other_max = float other_max = XLLM_SHFL_XOR_SYNC_WIDTH(
XLLM_SHFL_XOR_SYNC_WIDTH(0xffffffff, max_val, mask, THREADS_PER_ROW); kSoftmaxFullMask, max_val, mask, kThreadsPerRow);
int other_expert = int other_expert = XLLM_SHFL_XOR_SYNC_WIDTH(
XLLM_SHFL_XOR_SYNC_WIDTH(0xffffffff, expert, mask, THREADS_PER_ROW); kSoftmaxFullMask, expert, mask, kThreadsPerRow);
// We want lower indices to "win" in every thread so we break ties this // We want lower indices to "win" in every thread so we break ties this
// way // 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 // Finally, we clear the value in the thread with the current max if there
// is another iteration to run. // is another iteration to run.
if (k_idx + 1 < k) { 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 = 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 // Only the thread in the group which produced the max will reset the
// "winning" value to -inf. // "winning" value to -inf.
if (thread_group_idx == thread_to_clear_in_group) { 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 // Safe to set to any negative value since row_chunk values must be
// between 0 and 1. // 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; -10000.f;
} }
} }
@@ -623,18 +635,17 @@ void topk_gating_softmax_launcher_helper(const T* input,
const float moe_softcapping, const float moe_softcapping,
const float* correction_bias, const float* correction_bias,
cudaStream_t stream) { 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 = static constexpr int kBytesPerLdg = MIN(kMaxBytesPerLdg, sizeof(T) * EXPERTS);
MIN(MAX_BYTES_PER_LDG, sizeof(T) * EXPERTS); using Constants = TopkConstants<T, EXPERTS, kBytesPerLdg>;
using Constants = TopkConstants<T, EXPERTS, BYTES_PER_LDG>; static constexpr int kVpt = Constants::VPT;
static constexpr int VPT = Constants::VPT; static constexpr int kRowsPerWarp = Constants::ROWS_PER_WARP;
static constexpr int ROWS_PER_WARP = Constants::ROWS_PER_WARP; const int num_warps = (num_rows + kRowsPerWarp - 1) / kRowsPerWarp;
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; const int num_blocks = (num_warps + WARPS_PER_TB - 1) / WARPS_PER_TB;
dim3 block_dim(WARP_SIZE, 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, <<<num_blocks, block_dim, 0, stream>>>(input,
finished, finished,
output, output,
@@ -675,69 +686,69 @@ void topk_gating_softmax_kernel_launcher(const T* gating_output,
const float moe_softcapping, const float moe_softcapping,
const float* correction_bias, const float* correction_bias,
cudaStream_t stream) { cudaStream_t stream) {
static constexpr int WARPS_PER_TB = 4; static constexpr int kWarpsPerTb = 4;
switch (num_experts) { switch (num_experts) {
case 1: case 1:
LAUNCH_SOFTMAX(T, 1, WARPS_PER_TB); LAUNCH_SOFTMAX(T, 1, kWarpsPerTb);
break; break;
case 2: case 2:
LAUNCH_SOFTMAX(T, 2, WARPS_PER_TB); LAUNCH_SOFTMAX(T, 2, kWarpsPerTb);
break; break;
case 4: case 4:
LAUNCH_SOFTMAX(T, 4, WARPS_PER_TB); LAUNCH_SOFTMAX(T, 4, kWarpsPerTb);
break; break;
case 8: case 8:
LAUNCH_SOFTMAX(T, 8, WARPS_PER_TB); LAUNCH_SOFTMAX(T, 8, kWarpsPerTb);
break; break;
case 16: case 16:
LAUNCH_SOFTMAX(T, 16, WARPS_PER_TB); LAUNCH_SOFTMAX(T, 16, kWarpsPerTb);
break; break;
case 32: case 32:
LAUNCH_SOFTMAX(T, 32, WARPS_PER_TB); LAUNCH_SOFTMAX(T, 32, kWarpsPerTb);
break; break;
case 64: case 64:
LAUNCH_SOFTMAX(T, 64, WARPS_PER_TB); LAUNCH_SOFTMAX(T, 64, kWarpsPerTb);
break; break;
case 128: case 128:
LAUNCH_SOFTMAX(T, 128, WARPS_PER_TB); LAUNCH_SOFTMAX(T, 128, kWarpsPerTb);
break; break;
case 256: case 256:
LAUNCH_SOFTMAX(T, 256, WARPS_PER_TB); LAUNCH_SOFTMAX(T, 256, kWarpsPerTb);
break; break;
default: { default: {
CHECK(softmax_workspace != nullptr) CHECK(softmax_workspace != nullptr)
<< "softmax_workspace must be provided for num_experts that are " << "softmax_workspace must be provided for num_experts that are "
"not a power of 2."; "not a power of 2.";
static constexpr int TPB = 256; static constexpr int kTpb = 256;
moe_softmax<T, TPB><<<num_tokens, TPB, 0, stream>>>(gating_output, moe_softmax<T, kTpb><<<num_tokens, kTpb, 0, stream>>>(gating_output,
nullptr, nullptr,
softmax_workspace, softmax_workspace,
num_experts, num_experts,
moe_softcapping, moe_softcapping,
correction_bias); correction_bias);
if (topk == 1) { if (topk == 1) {
// Note: As an optimization for better performance, // Note: As an optimization for better performance,
// the softmax_workspace is overwritten in-place by both moeTopK and // the softmax_workspace is overwritten in-place by both moeTopK and
// moe_topk_fast. // moe_topk_fast.
moe_topK<TPB><<<num_tokens, TPB, 0, stream>>>(softmax_workspace, moe_topK<kTpb><<<num_tokens, kTpb, 0, stream>>>(softmax_workspace,
nullptr, nullptr,
topk_weights, topk_weights,
topk_indices, topk_indices,
num_experts, num_experts,
topk, topk,
0, 0,
num_experts, num_experts,
renormalize); renormalize);
} else { } else {
moe_topk_fast<TPB><<<num_tokens, TPB, 0, stream>>>(softmax_workspace, moe_topk_fast<kTpb><<<num_tokens, kTpb, 0, stream>>>(softmax_workspace,
nullptr, nullptr,
topk_weights, topk_weights,
topk_indices, topk_indices,
num_experts, num_experts,
topk, topk,
0, 0,
num_experts, num_experts,
renormalize); renormalize);
} }
} }
} }
@@ -835,8 +846,8 @@ void topk_softmax(torch::Tensor& topk_weights, // [num_tokens, topk]
bias_ptr, bias_ptr,
stream); stream);
} else if (dtype == at::ScalarType::BFloat16) { } else if (dtype == at::ScalarType::BFloat16) {
topk_gating_softmax_kernel_launcher<__nv_bfloat16>( topk_gating_softmax_kernel_launcher<BFloat16Type>(
reinterpret_cast<const __nv_bfloat16*>( reinterpret_cast<const BFloat16Type*>(
gating_output.data_ptr<at::BFloat16>()), gating_output.data_ptr<at::BFloat16>()),
topk_weights.data_ptr<float>(), topk_weights.data_ptr<float>(),
topk_indices.data_ptr<int>(), 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 // Rejection sampling parameters for speculative decoding
// PRD #0 (9e0f1402): Added mode field for explicit greedy/probabilistic dispatch
struct RejectionSampleParams { struct RejectionSampleParams {
// Candidate draft token indices to be verified. // Candidate draft token indices to be verified.
// Shape: [total_draft_tokens]. Dtype: int32. // Shape: [total_draft_tokens]. Dtype: int32.
@@ -849,6 +850,10 @@ struct RejectionSampleParams {
// The maximum number of draft tokens in the batch (max value in // The maximum number of draft tokens in the batch (max value in
// num_draft_tokens). // num_draft_tokens).
int32_t max_spec_len; 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 // 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. // Whether to apply L2 norm to q and k inside the kernel. Default: false.
bool use_qk_l2norm_in_kernel = 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/DynamicLibrary.h>
#include <ATen/core/dispatch/Dispatcher.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 <optional>
#include <tuple> #include <tuple>
@@ -54,7 +55,7 @@ void block_copy(torch::Tensor key_cache_ptrs,
torch::Tensor cum_sum, torch::Tensor cum_sum,
int64_t numel_per_block, int64_t numel_per_block,
torch::ScalarType cache_dtype); torch::ScalarType cache_dtype);
#if !defined(USE_DCU) #if !defined(USE_DCU) && !defined(__ILUVATAR__)
void batch_prefill(const std::string& uri, void batch_prefill(const std::string& uri,
ffi::Array<int64_t> plan_info, ffi::Array<int64_t> plan_info,
torch::Tensor float_workspace_buffer, torch::Tensor float_workspace_buffer,
@@ -142,7 +143,7 @@ void batch_decode(const std::string& uri,
std::optional<torch::Tensor>& output_lse, std::optional<torch::Tensor>& output_lse,
bool use_tensor_core, bool use_tensor_core,
std::optional<torch::Tensor> qo_indptr = std::nullopt); std::optional<torch::Tensor> qo_indptr = std::nullopt);
#endif // !defined(USE_DCU) #endif // !defined(USE_DCU) && !defined(__ILUVATAR__)
void rms_norm(torch::Tensor output, void rms_norm(torch::Tensor output,
torch::Tensor input, torch::Tensor input,
torch::Tensor weight, torch::Tensor weight,
@@ -303,4 +304,4 @@ torch::Tensor moe_combine_result(const torch::Tensor& gemm2,
int64_t N, int64_t N,
int32_t topk); 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)) || \ #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 // CUDA < 12.0 runs into issues with packed type conversion
template <> template <>
class _typeConvert<c10::Half> { 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 #else
#include <c10/cuda/CUDAGuard.h> #include <c10/cuda/CUDAGuard.h>
#endif #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> #include <torch/torch.h>
#if !defined(USE_DCU) #if !defined(USE_DCU) && !defined(__ILUVATAR__)
#include <tvm/ffi/container/array.h> #include <tvm/ffi/container/array.h>
#include <tvm/ffi/container/tensor.h> #include <tvm/ffi/container/tensor.h>
#include <tvm/ffi/extra/c_env_api.h> #include <tvm/ffi/extra/c_env_api.h>
@@ -46,7 +47,7 @@ limitations under the License.
#define HOST_INLINE inline #define HOST_INLINE inline
#endif #endif
#if !defined(USE_DCU) #if !defined(USE_DCU) && !defined(__ILUVATAR__)
namespace ffi = tvm::ffi; namespace ffi = tvm::ffi;
#endif #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); 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); DLDataType to_dl_data_type(torch::ScalarType scalar_type);
// below are tvm-ffi related functions // below are tvm-ffi related functions
@@ -159,5 +160,5 @@ inline void bind_tvmffi_stream_to_current_torch_stream(
<< " dev=" << device.index(); << " dev=" << device.index();
} }
} }
#endif // !defined(USE_DCU) #endif // !defined(USE_DCU) && !defined(__ILUVATAR__)
} // namespace xllm::kernel::cuda } // 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. limitations under the License.
==============================================================================*/ ==============================================================================*/
#include "kernels/cuda/cuda_ops_api.h" #include "cuda_ops_api.h"
#include "kernels/cuda/utils.h" #include "utils.h"
#include "platform/device.h"
#include "platform/platform.h"
namespace xllm::kernel::cuda { namespace xllm::kernel::cuda {
@@ -51,74 +51,73 @@ torch::Tensor cutlass_fused_moe(
bool use_packed_weights, bool use_packed_weights,
int32_t tune_max_num_tokens, int32_t tune_max_num_tokens,
ActivationType activation_type) { 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 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) { if (min_latency_mode) {
num_rows *= fc2_expert_weights.size(0); num_rows *= fc2_expert_weights.size(0);
} }
std::vector<int64_t> output_shape = {num_rows, hidden_size};
torch::Tensor result_output; torch::Tensor result_output;
if (output.has_value() && output.value().defined()) { if (output.has_value() && output.value().defined()) {
result_output = output.value(); result_output = output.value();
} else { } else {
torch::TensorOptions options = input.options().dtype(output_dtype); result_output = torch::zeros({num_rows, hidden_size},
result_output = torch::empty(output_shape, options); input.options().dtype(output_dtype));
} }
std::string fused_moe_uri = "fused_moe"; if (Platform::is_support_ivcore10()) {
if (Platform::is_support_sm90a()) { // BI-V100 path: per-token expert gather + matmul + SiLU-gate + matmul
fused_moe_uri += "_90"; // This replaces the tvm ffi CUTLASS path with native PyTorch ops.
} else if (Platform::is_support_sm100a() || Platform::is_support_sm100f()) { for (int64_t t = 0; t < num_tokens; ++t) {
fused_moe_uri += "_100"; auto token = input[t].unsqueeze(0); // [1, hidden]
} else if (Platform::is_support_sm120a()) { torch::Tensor accum = torch::zeros({1, hidden_size},
fused_moe_uri += "_120"; input.options().dtype(output_dtype));
} else { for (int64_t k = 0; k < top_k; ++k) {
LOG(FATAL) << "FusedMoE is only supported on sm90, sm100, sm120."; 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()); // Original NVIDIA GPU path (sm90/sm100/sm120) via tvm ffi
TORCH_CHECK(false,
ffi::Module fused_moe_runner = "cutlass_fused_moe: no supported platform. "
get_function(fused_moe_uri, "init")( "BI-V100 should use ivcore10 path above; "
to_dl_data_type(input.scalar_type()), "NVIDIA GPUs require sm90+.");
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; 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> #include <cub/util_type.cuh>
#if !defined(USE_DCU) && !defined(USE_MACA) // <cuda/functional> requires CUDA 12+ (libcudacxx); BI-V100 runs CUDA 10.2
#endif // and does not ship that header. The include is unused in this file anyway.
#include "device_utils.cuh" #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 // Rejection sampling parameters for speculative decoding
// PRD #0 (9e0f1402): Added mode field for explicit greedy/probabilistic dispatch
struct RejectionSampleParams { struct RejectionSampleParams {
// Candidate draft token indices to be verified. // Candidate draft token indices to be verified.
// Shape: [total_draft_tokens]. Dtype: int32. // Shape: [total_draft_tokens]. Dtype: int32.
@@ -849,6 +850,10 @@ struct RejectionSampleParams {
// The maximum number of draft tokens in the batch (max value in // The maximum number of draft tokens in the batch (max value in
// num_draft_tokens). // num_draft_tokens).
int32_t max_spec_len; 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 // 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. // Whether to apply L2 norm to q and k inside the kernel. Default: false.
bool use_qk_l2norm_in_kernel = false; bool use_qk_l2norm_in_kernel = false;
}; };
} // namespace xllm::kernel } // namespace xllm::kernel

157
qwen3_6_scripts/moe_topk_sigmoid_kernels.cuh
View File

@@ -23,12 +23,21 @@ limitations under the License.
#include <cub/util_type.cuh> #include <cub/util_type.cuh>
#if !defined(USE_DCU) && !defined(USE_MACA)
#endif
#include "device_utils.cuh" #include "device_utils.cuh"
namespace { namespace {
using namespace xllm::kernel::cuda; 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 =============================== // ====================== Sigmoid things ===============================
// We have our own implementation of sigmoid here so we can support transposing // 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 // the output in the sigmoid kernel when we extend this module to support
@@ -182,29 +191,29 @@ __launch_bounds__(WARPS_PER_CTA* WARP_SIZE) __global__
static_assert(BYTES_PER_LDG <= 16, "BYTES_PER_LDG must be leq 16"); static_assert(BYTES_PER_LDG <= 16, "BYTES_PER_LDG must be leq 16");
// Number of bytes each thread pulls in per load // Number of bytes each thread pulls in per load
static constexpr int ELTS_PER_LDG = BYTES_PER_LDG / sizeof(T); static constexpr int kEltsPerLdg = BYTES_PER_LDG / sizeof(T);
static constexpr int ELTS_PER_ROW = NUM_EXPERTS; static constexpr int kEltsPerRow = NUM_EXPERTS;
static constexpr int THREADS_PER_ROW = ELTS_PER_ROW / VPT; static constexpr int kThreadsPerRow = kEltsPerRow / VPT;
static constexpr int LDG_PER_THREAD = VPT / ELTS_PER_LDG; static constexpr int kLdgPerThread = VPT / kEltsPerLdg;
// Restrictions based on previous section. // Restrictions based on previous section.
static_assert( static_assert(
VPT % ELTS_PER_LDG == 0, VPT % kEltsPerLdg == 0,
"The elements per thread must be a multiple of the elements per ldg"); "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"); "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"); "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"); "THREADS_PER_ROW can be at most warp size");
// We have NUM_EXPERTS elements per row. We specialize for small #experts // We have NUM_EXPERTS elements per row. We specialize for small #experts
static constexpr int ELTS_PER_WARP = WARP_SIZE * VPT; static constexpr int kEltsPerWarp = WARP_SIZE * VPT;
static constexpr int ROWS_PER_WARP = ELTS_PER_WARP / ELTS_PER_ROW; static constexpr int kRowsPerWarp = kEltsPerWarp / kEltsPerRow;
static constexpr int ROWS_PER_CTA = WARPS_PER_CTA * ROWS_PER_WARP; static constexpr int kRowsPerCta = WARPS_PER_CTA * kRowsPerWarp;
// Restrictions for previous section. // 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"); "The elts per row must cleanly divide the total elt per warp");
// ===================== From this point, we finally start computing run-time // ===================== From this point, we finally start computing run-time
@@ -213,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 // 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 // block contains WARPS_PER_CTA warps. This, each block processes a chunk of
// rows. We start by computing the start row for each block. // 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. // 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. // 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 // 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; const int thread_row = warp_base_row + thread_row_in_warp;
// Threads with indices out of bounds should early exit here. // Threads with indices out of bounds should early exit here.
@@ -231,12 +240,12 @@ __launch_bounds__(WARPS_PER_CTA* WARP_SIZE) __global__
// We finally start setting up the read pointers for each thread. First, each // We finally start setting up the read pointers for each thread. First, each
// thread jumps to the start of the row it will read. // 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 // Now, we compute the group each thread belong to in order to determine the
// first column to start loads. // first column to start loads.
const int thread_group_idx = threadIdx.x % THREADS_PER_ROW; const int thread_group_idx = threadIdx.x % kThreadsPerRow;
const int first_elt_read_by_thread = thread_group_idx * ELTS_PER_LDG; const int first_elt_read_by_thread = thread_group_idx * kEltsPerLdg;
const T* thread_read_ptr = thread_row_ptr + first_elt_read_by_thread; 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 // Determine the pointer type to use to read in the data depending on the
@@ -244,7 +253,7 @@ __launch_bounds__(WARPS_PER_CTA* WARP_SIZE) __global__
// up to 16. NOTE(woosuk): The original implementation uses CUTLASS aligned // 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 // array here. We defined our own aligned array and use it here to avoid the
// dependency on CUTLASS. // dependency on CUTLASS.
using AccessType = AlignedArray<T, ELTS_PER_LDG>; using AccessType = AlignedArray<T, kEltsPerLdg>;
// Finally, we pull in the data from global mem // Finally, we pull in the data from global mem
T row_chunk_temp[VPT]; T row_chunk_temp[VPT];
@@ -256,8 +265,8 @@ __launch_bounds__(WARPS_PER_CTA* WARP_SIZE) __global__
// Note(Byron): interleaved loads to achieve better memory coalescing // Note(Byron): interleaved loads to achieve better memory coalescing
// | thread[0] | thread[1] | thread[2] | thread[3] | thread[0] | thread[1] | // | thread[0] | thread[1] | thread[2] | thread[3] | thread[0] | thread[1] |
// thread[2] | thread[3] | ... // thread[2] | thread[3] | ...
for (int ii = 0; ii < LDG_PER_THREAD; ++ii) { for (int ii = 0; ii < kLdgPerThread; ++ii) {
row_chunk_vec_ptr[ii] = vec_thread_read_ptr[ii * THREADS_PER_ROW]; row_chunk_vec_ptr[ii] = vec_thread_read_ptr[ii * kThreadsPerRow];
} }
float row_chunk[VPT]; float row_chunk[VPT];
@@ -274,11 +283,10 @@ __launch_bounds__(WARPS_PER_CTA* WARP_SIZE) __global__
|--------- group0 --------| |----------group1 --------| |--------- group0 --------| |----------group1 --------|
^ local2 ^ local2
*/ */
const int group_id = ii / ELTS_PER_LDG; const int group_id = ii / kEltsPerLdg;
const int local_id = ii % ELTS_PER_LDG; const int local_id = ii % kEltsPerLdg;
const int expert_idx = first_elt_read_by_thread + const int expert_idx = first_elt_read_by_thread +
group_id * THREADS_PER_ROW * ELTS_PER_LDG + group_id * kThreadsPerRow * kEltsPerLdg + local_id;
local_id;
val = val + correction_bias[expert_idx]; val = val + correction_bias[expert_idx];
} }
@@ -288,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 // 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. // the topk elements in each row, along with the max index.
int start_col = first_elt_read_by_thread; 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; float row_sum_for_renormalize = 0;
@@ -297,11 +305,11 @@ __launch_bounds__(WARPS_PER_CTA* WARP_SIZE) __global__
float max_val = row_chunk[0]; float max_val = row_chunk[0];
int expert = start_col; int expert = start_col;
#pragma unroll #pragma unroll
for (int ldg = 0, col = start_col; ldg < LDG_PER_THREAD; for (int ldg = 0, col = start_col; ldg < kLdgPerThread;
++ldg, col += COLS_PER_GROUP_LDG) { ++ldg, col += kColsPerGroupLdg) {
#pragma unroll #pragma unroll
for (int ii = 0; ii < ELTS_PER_LDG; ++ii) { for (int ii = 0; ii < kEltsPerLdg; ++ii) {
float val = row_chunk[ldg * ELTS_PER_LDG + ii]; float val = row_chunk[ldg * kEltsPerLdg + ii];
// No check on the experts here since columns with the smallest index // No check on the experts here since columns with the smallest index
// are processed first and only updated if > (not >=) // are processed first and only updated if > (not >=)
@@ -317,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 // 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... // their max with -inf and the warp can run more iterations...
#pragma unroll #pragma unroll
for (int mask = THREADS_PER_ROW / 2; mask > 0; mask /= 2) { for (int mask = kThreadsPerRow / 2; mask > 0; mask /= 2) {
float other_max = float other_max = XLLM_SHFL_XOR_SYNC_WIDTH(
XLLM_SHFL_XOR_SYNC_WIDTH(0xffffffff, max_val, mask, THREADS_PER_ROW); kSigmoidFullMask, max_val, mask, kThreadsPerRow);
int other_expert = int other_expert = XLLM_SHFL_XOR_SYNC_WIDTH(
XLLM_SHFL_XOR_SYNC_WIDTH(0xffffffff, expert, mask, THREADS_PER_ROW); kSigmoidFullMask, expert, mask, kThreadsPerRow);
// We want lower indices to "win" in every thread so we break ties this // We want lower indices to "win" in every thread so we break ties this
// way // way
@@ -354,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 // Finally, we clear the value in the thread with the current max if there
// is another iteration to run. // is another iteration to run.
if (k_idx + 1 < k) { 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 = 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 // Only the thread in the group which produced the max will reset the
// "winning" value to -inf. // "winning" value to -inf.
if (thread_group_idx == thread_to_clear_in_group) { 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 // Safe to set to any negative value since row_chunk values must be
// between 0 and 1. // 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; -10000.f;
} }
} }
@@ -393,18 +401,17 @@ void topk_gating_sigmoid_launcher_helper(const T* input,
const bool renormalize, const bool renormalize,
const float* correction_bias, const float* correction_bias,
cudaStream_t stream) { 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 = static constexpr int kBytesPerLdg = MIN(kMaxBytesPerLdg, sizeof(T) * EXPERTS);
MIN(MAX_BYTES_PER_LDG, sizeof(T) * EXPERTS); using Constants = TopkConstants<T, EXPERTS, kBytesPerLdg>;
using Constants = TopkConstants<T, EXPERTS, BYTES_PER_LDG>; static constexpr int kVpt = Constants::VPT;
static constexpr int VPT = Constants::VPT; static constexpr int kRowsPerWarp = Constants::ROWS_PER_WARP;
static constexpr int ROWS_PER_WARP = Constants::ROWS_PER_WARP; const int num_warps = (num_rows + kRowsPerWarp - 1) / kRowsPerWarp;
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; const int num_blocks = (num_warps + WARPS_PER_TB - 1) / WARPS_PER_TB;
dim3 block_dim(WARP_SIZE, 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, <<<num_blocks, block_dim, 0, stream>>>(input,
finished, finished,
output, output,
@@ -442,55 +449,55 @@ void topk_gating_sigmoid_kernel_launcher(const T* gating_output,
const bool renormalize, const bool renormalize,
const float* correction_bias, const float* correction_bias,
cudaStream_t stream) { cudaStream_t stream) {
static constexpr int WARPS_PER_TB = 4; static constexpr int kWarpsPerTb = 4;
switch (num_experts) { switch (num_experts) {
case 1: case 1:
LAUNCH_SIGMOID(T, 1, WARPS_PER_TB); LAUNCH_SIGMOID(T, 1, kWarpsPerTb);
break; break;
case 2: case 2:
LAUNCH_SIGMOID(T, 2, WARPS_PER_TB); LAUNCH_SIGMOID(T, 2, kWarpsPerTb);
break; break;
case 4: case 4:
LAUNCH_SIGMOID(T, 4, WARPS_PER_TB); LAUNCH_SIGMOID(T, 4, kWarpsPerTb);
break; break;
case 8: case 8:
LAUNCH_SIGMOID(T, 8, WARPS_PER_TB); LAUNCH_SIGMOID(T, 8, kWarpsPerTb);
break; break;
case 16: case 16:
LAUNCH_SIGMOID(T, 16, WARPS_PER_TB); LAUNCH_SIGMOID(T, 16, kWarpsPerTb);
break; break;
case 32: case 32:
LAUNCH_SIGMOID(T, 32, WARPS_PER_TB); LAUNCH_SIGMOID(T, 32, kWarpsPerTb);
break; break;
case 64: case 64:
LAUNCH_SIGMOID(T, 64, WARPS_PER_TB); LAUNCH_SIGMOID(T, 64, kWarpsPerTb);
break; break;
case 128: case 128:
LAUNCH_SIGMOID(T, 128, WARPS_PER_TB); LAUNCH_SIGMOID(T, 128, kWarpsPerTb);
break; break;
case 256: case 256:
LAUNCH_SIGMOID(T, 256, WARPS_PER_TB); LAUNCH_SIGMOID(T, 256, kWarpsPerTb);
break; break;
default: { default: {
TORCH_CHECK(sigmoid_workspace != nullptr, TORCH_CHECK(sigmoid_workspace != nullptr,
"sigmoid_workspace must be provided for num_experts that are " "sigmoid_workspace must be provided for num_experts that are "
"not a power of 2."); "not a power of 2.");
static constexpr int TPB = 256; static constexpr int kTpb = 256;
moe_sigmoid<T, TPB><<<num_tokens, TPB, 0, stream>>>(gating_output, moe_sigmoid<T, kTpb><<<num_tokens, kTpb, 0, stream>>>(gating_output,
nullptr, nullptr,
sigmoid_workspace, sigmoid_workspace,
num_experts, num_experts,
correction_bias); correction_bias);
moe_topK<TPB><<<num_tokens, TPB, 0, stream>>>(sigmoid_workspace, moe_topK<kTpb><<<num_tokens, kTpb, 0, stream>>>(sigmoid_workspace,
nullptr, nullptr,
topk_weights, topk_weights,
topk_indices, topk_indices,
num_experts, num_experts,
topk, topk,
0, 0,
num_experts, num_experts,
renormalize, renormalize,
correction_bias); correction_bias);
} }
} }
} }
@@ -582,8 +589,8 @@ void topk_sigmoid(torch::Tensor& topk_weights, // [num_tokens, topk]
bias_ptr, bias_ptr,
stream); stream);
} else if (dtype == at::ScalarType::BFloat16) { } else if (dtype == at::ScalarType::BFloat16) {
topk_gating_sigmoid_kernel_launcher<__nv_bfloat16>( topk_gating_sigmoid_kernel_launcher<BFloat16Type>(
reinterpret_cast<const __nv_bfloat16*>( reinterpret_cast<const BFloat16Type*>(
gating_output.data_ptr<at::BFloat16>()), gating_output.data_ptr<at::BFloat16>()),
topk_weights.data_ptr<float>(), topk_weights.data_ptr<float>(),
topk_indices.data_ptr<int>(), topk_indices.data_ptr<int>(),

246
qwen3_6_scripts/moe_topk_softmax_kernels.cuh
View File

@@ -23,6 +23,9 @@ limitations under the License.
#include <cub/util_type.cuh> #include <cub/util_type.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" #include "device_utils.cuh"
using cub_kvp = cub::KeyValuePair<int, float>; using cub_kvp = cub::KeyValuePair<int, float>;
@@ -31,6 +34,12 @@ namespace {
using namespace xllm::kernel::cuda; 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 =============================== // ====================== Softmax things ===============================
// We have our own implementation of softmax here so we can support transposing // 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 // the output in the softmax kernel when we extend this module to support
@@ -110,9 +119,10 @@ __launch_bounds__(TPB) __global__
} }
namespace moe { namespace moe {
struct TopKPair { class TopKPair {
static const int PAIR = 2; public:
static const int MAX_INDEX = 0; static constexpr int kPair = 2;
static constexpr int kMaxIndex = 0;
cub_kvp max; cub_kvp max;
cub_kvp secondMax; cub_kvp secondMax;
@@ -121,7 +131,8 @@ struct TopKPair {
: max(max), secondMax(secondMax) {} : max(max), secondMax(secondMax) {}
}; };
struct TopKPairArgMax { class TopKPairArgMax {
public:
__device__ TopKPairArgMax() {} __device__ TopKPairArgMax() {}
__device__ __forceinline__ TopKPair __device__ __forceinline__ TopKPair
operator()(const TopKPair& candidate1, const TopKPair& candidate2) const { operator()(const TopKPair& candidate1, const TopKPair& candidate2) const {
@@ -175,8 +186,8 @@ __launch_bounds__(TPB) __global__
const int thread_read_offset = blockIdx.x * num_experts; const int thread_read_offset = blockIdx.x * num_experts;
float row_sum_for_renormalize = 0; float row_sum_for_renormalize = 0;
// Each loop finds the top 2 elements, // Each loop finds the top 2 elements,
// thus requiring only ⌈k/2⌉ loops (calculated as (k + 1) / 2). // thus requiring only ceil(k / 2) loops (calculated as (k + 1) / 2).
for (int k_idx = 0; k_idx < (k + TopKPair::PAIR - 1) / TopKPair::PAIR; for (int k_idx = 0; k_idx < (k + TopKPair::kPair - 1) / TopKPair::kPair;
++k_idx) { ++k_idx) {
// Initializing the top 2 elements by the minimum value. // Initializing the top 2 elements by the minimum value.
thread_pair.max.key = 0; thread_pair.max.key = 0;
@@ -204,9 +215,11 @@ __launch_bounds__(TPB) __global__
if (threadIdx.x == 0) { if (threadIdx.x == 0) {
#pragma unroll #pragma unroll
// updating 2 elements to the result. // updating 2 elements to the result.
for (int i = 0; i < TopKPair::PAIR; i++) { for (int i = 0; i < TopKPair::kPair; i++) {
if (k_idx * 2 + i >= k) break; if (k_idx * 2 + i >= k) {
cub_kvp result = (i == TopKPair::MAX_INDEX) ? result_pair.max break;
}
cub_kvp result = (i == TopKPair::kMaxIndex) ? result_pair.max
: result_pair.secondMax; : result_pair.secondMax;
int expert = result.key; int expert = result.key;
bool node_uses_expert = expert >= start_expert && expert < end_expert; bool node_uses_expert = expert >= start_expert && expert < end_expert;
@@ -342,29 +355,29 @@ __launch_bounds__(WARPS_PER_CTA* WARP_SIZE) __global__
static_assert(BYTES_PER_LDG <= 16, "BYTES_PER_LDG must be leq 16"); static_assert(BYTES_PER_LDG <= 16, "BYTES_PER_LDG must be leq 16");
// Number of bytes each thread pulls in per load // Number of bytes each thread pulls in per load
static constexpr int ELTS_PER_LDG = BYTES_PER_LDG / sizeof(T); static constexpr int kEltsPerLdg = BYTES_PER_LDG / sizeof(T);
static constexpr int ELTS_PER_ROW = NUM_EXPERTS; static constexpr int kEltsPerRow = NUM_EXPERTS;
static constexpr int THREADS_PER_ROW = ELTS_PER_ROW / VPT; static constexpr int kThreadsPerRow = kEltsPerRow / VPT;
static constexpr int LDG_PER_THREAD = VPT / ELTS_PER_LDG; static constexpr int kLdgPerThread = VPT / kEltsPerLdg;
// Restrictions based on previous section. // Restrictions based on previous section.
static_assert( static_assert(
VPT % ELTS_PER_LDG == 0, VPT % kEltsPerLdg == 0,
"The elements per thread must be a multiple of the elements per ldg"); "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"); "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"); "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"); "THREADS_PER_ROW can be at most warp size");
// We have NUM_EXPERTS elements per row. We specialize for small #experts // We have NUM_EXPERTS elements per row. We specialize for small #experts
static constexpr int ELTS_PER_WARP = WARP_SIZE * VPT; static constexpr int kEltsPerWarp = WARP_SIZE * VPT;
static constexpr int ROWS_PER_WARP = ELTS_PER_WARP / ELTS_PER_ROW; static constexpr int kRowsPerWarp = kEltsPerWarp / kEltsPerRow;
static constexpr int ROWS_PER_CTA = WARPS_PER_CTA * ROWS_PER_WARP; static constexpr int kRowsPerCta = WARPS_PER_CTA * kRowsPerWarp;
// Restrictions for previous section. // 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"); "The elts per row must cleanly divide the total elt per warp");
// ===================== From this point, we finally start computing run-time // ===================== From this point, we finally start computing run-time
@@ -373,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 // 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 // block contains WARPS_PER_CTA warps. This, each block processes a chunk of
// rows. We start by computing the start row for each block. // 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. // 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. // 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 // 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; const int thread_row = warp_base_row + thread_row_in_warp;
// Threads with indices out of bounds should early exit here. // Threads with indices out of bounds should early exit here.
@@ -391,12 +404,12 @@ __launch_bounds__(WARPS_PER_CTA* WARP_SIZE) __global__
// We finally start setting up the read pointers for each thread. First, each // We finally start setting up the read pointers for each thread. First, each
// thread jumps to the start of the row it will read. // 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 // Now, we compute the group each thread belong to in order to determine the
// first column to start loads. // first column to start loads.
const int thread_group_idx = threadIdx.x % THREADS_PER_ROW; const int thread_group_idx = threadIdx.x % kThreadsPerRow;
const int first_elt_read_by_thread = thread_group_idx * ELTS_PER_LDG; const int first_elt_read_by_thread = thread_group_idx * kEltsPerLdg;
const T* thread_read_ptr = thread_row_ptr + first_elt_read_by_thread; 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 // Determine the pointer type to use to read in the data depending on the
@@ -404,7 +417,7 @@ __launch_bounds__(WARPS_PER_CTA* WARP_SIZE) __global__
// up to 16. NOTE(woosuk): The original implementation uses CUTLASS aligned // 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 // array here. We defined our own aligned array and use it here to avoid the
// dependency on CUTLASS. // dependency on CUTLASS.
using AccessType = AlignedArray<T, ELTS_PER_LDG>; using AccessType = AlignedArray<T, kEltsPerLdg>;
// Finally, we pull in the data from global mem // Finally, we pull in the data from global mem
T row_chunk_temp[VPT]; T row_chunk_temp[VPT];
@@ -416,8 +429,8 @@ __launch_bounds__(WARPS_PER_CTA* WARP_SIZE) __global__
// Note(Byron): interleaved loads to achieve better memory coalescing // Note(Byron): interleaved loads to achieve better memory coalescing
// | thread[0] | thread[1] | thread[2] | thread[3] | thread[0] | thread[1] | // | thread[0] | thread[1] | thread[2] | thread[3] | thread[0] | thread[1] |
// thread[2] | thread[3] | ... // thread[2] | thread[3] | ...
for (int ii = 0; ii < LDG_PER_THREAD; ++ii) { for (int ii = 0; ii < kLdgPerThread; ++ii) {
row_chunk_vec_ptr[ii] = vec_thread_read_ptr[ii * THREADS_PER_ROW]; row_chunk_vec_ptr[ii] = vec_thread_read_ptr[ii * kThreadsPerRow];
} }
float row_chunk[VPT]; float row_chunk[VPT];
@@ -446,10 +459,10 @@ __launch_bounds__(WARPS_PER_CTA* WARP_SIZE) __global__
|--------- group0 --------| |----------group1 --------| |--------- group0 --------| |----------group1 --------|
^ local2 ^ local2
*/ */
const int group_id = ii / ELTS_PER_LDG; const int group_id = ii / kEltsPerLdg;
const int local_id = ii % ELTS_PER_LDG; const int local_id = ii % kEltsPerLdg;
const int expert_idx = first_elt_read_by_thread + const int expert_idx = first_elt_read_by_thread +
group_id * THREADS_PER_ROW * ELTS_PER_LDG + group_id * kThreadsPerRow * kEltsPerLdg +
local_id; local_id;
val = val + correction_bias[expert_idx]; val = val + correction_bias[expert_idx];
} }
@@ -474,11 +487,11 @@ __launch_bounds__(WARPS_PER_CTA* WARP_SIZE) __global__
// Now, we find the max within the thread group and distribute among the // 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 // threads. We use a butterfly reduce. lane id: 0-31 within a warp
#pragma unroll #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) // butterfly reduce with (lane id ^ mask)
thread_max = max(thread_max, thread_max = max(thread_max,
XLLM_SHFL_XOR_SYNC_WIDTH( 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. // From this point, thread max in all the threads have the max within the row.
@@ -494,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 // Now, we perform the sum reduce within each thread group. Similar to the max
// reduce, we use a bufferfly pattern. // reduce, we use a bufferfly pattern.
#pragma unroll #pragma unroll
for (int mask = THREADS_PER_ROW / 2; mask > 0; mask /= 2) { for (int mask = kThreadsPerRow / 2; mask > 0; mask /= 2) {
row_sum += row_sum += XLLM_SHFL_XOR_SYNC_WIDTH(
XLLM_SHFL_XOR_SYNC_WIDTH(0xffffffff, row_sum, mask, THREADS_PER_ROW); kSoftmaxFullMask, row_sum, mask, kThreadsPerRow);
} }
// From this point, all threads have the max and the sum for their rows in the // From this point, all threads have the max and the sum for their rows in the
@@ -519,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 // 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. // the topk elements in each row, along with the max index.
int start_col = first_elt_read_by_thread; 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; float row_sum_for_renormalize = 0;
@@ -528,11 +541,11 @@ __launch_bounds__(WARPS_PER_CTA* WARP_SIZE) __global__
float max_val = row_chunk[0]; float max_val = row_chunk[0];
int expert = start_col; int expert = start_col;
#pragma unroll #pragma unroll
for (int ldg = 0, col = start_col; ldg < LDG_PER_THREAD; for (int ldg = 0, col = start_col; ldg < kLdgPerThread;
++ldg, col += COLS_PER_GROUP_LDG) { ++ldg, col += kColsPerGroupLdg) {
#pragma unroll #pragma unroll
for (int ii = 0; ii < ELTS_PER_LDG; ++ii) { for (int ii = 0; ii < kEltsPerLdg; ++ii) {
float val = row_chunk[ldg * ELTS_PER_LDG + ii]; float val = row_chunk[ldg * kEltsPerLdg + ii];
// No check on the experts here since columns with the smallest index // No check on the experts here since columns with the smallest index
// are processed first and only updated if > (not >=) // are processed first and only updated if > (not >=)
@@ -548,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 // 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... // their max with -inf and the warp can run more iterations...
#pragma unroll #pragma unroll
for (int mask = THREADS_PER_ROW / 2; mask > 0; mask /= 2) { for (int mask = kThreadsPerRow / 2; mask > 0; mask /= 2) {
float other_max = float other_max = XLLM_SHFL_XOR_SYNC_WIDTH(
XLLM_SHFL_XOR_SYNC_WIDTH(0xffffffff, max_val, mask, THREADS_PER_ROW); kSoftmaxFullMask, max_val, mask, kThreadsPerRow);
int other_expert = int other_expert = XLLM_SHFL_XOR_SYNC_WIDTH(
XLLM_SHFL_XOR_SYNC_WIDTH(0xffffffff, expert, mask, THREADS_PER_ROW); kSoftmaxFullMask, expert, mask, kThreadsPerRow);
// We want lower indices to "win" in every thread so we break ties this // We want lower indices to "win" in every thread so we break ties this
// way // way
@@ -582,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 // Finally, we clear the value in the thread with the current max if there
// is another iteration to run. // is another iteration to run.
if (k_idx + 1 < k) { 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 = 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 // Only the thread in the group which produced the max will reset the
// "winning" value to -inf. // "winning" value to -inf.
if (thread_group_idx == thread_to_clear_in_group) { 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 // Safe to set to any negative value since row_chunk values must be
// between 0 and 1. // 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; -10000.f;
} }
} }
@@ -622,18 +635,17 @@ void topk_gating_softmax_launcher_helper(const T* input,
const float moe_softcapping, const float moe_softcapping,
const float* correction_bias, const float* correction_bias,
cudaStream_t stream) { 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 = static constexpr int kBytesPerLdg = MIN(kMaxBytesPerLdg, sizeof(T) * EXPERTS);
MIN(MAX_BYTES_PER_LDG, sizeof(T) * EXPERTS); using Constants = TopkConstants<T, EXPERTS, kBytesPerLdg>;
using Constants = TopkConstants<T, EXPERTS, BYTES_PER_LDG>; static constexpr int kVpt = Constants::VPT;
static constexpr int VPT = Constants::VPT; static constexpr int kRowsPerWarp = Constants::ROWS_PER_WARP;
static constexpr int ROWS_PER_WARP = Constants::ROWS_PER_WARP; const int num_warps = (num_rows + kRowsPerWarp - 1) / kRowsPerWarp;
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; const int num_blocks = (num_warps + WARPS_PER_TB - 1) / WARPS_PER_TB;
dim3 block_dim(WARP_SIZE, 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, <<<num_blocks, block_dim, 0, stream>>>(input,
finished, finished,
output, output,
@@ -674,67 +686,69 @@ void topk_gating_softmax_kernel_launcher(const T* gating_output,
const float moe_softcapping, const float moe_softcapping,
const float* correction_bias, const float* correction_bias,
cudaStream_t stream) { cudaStream_t stream) {
static constexpr int WARPS_PER_TB = 4; static constexpr int kWarpsPerTb = 4;
switch (num_experts) { switch (num_experts) {
case 1: case 1:
LAUNCH_SOFTMAX(T, 1, WARPS_PER_TB); LAUNCH_SOFTMAX(T, 1, kWarpsPerTb);
break; break;
case 2: case 2:
LAUNCH_SOFTMAX(T, 2, WARPS_PER_TB); LAUNCH_SOFTMAX(T, 2, kWarpsPerTb);
break; break;
case 4: case 4:
LAUNCH_SOFTMAX(T, 4, WARPS_PER_TB); LAUNCH_SOFTMAX(T, 4, kWarpsPerTb);
break; break;
case 8: case 8:
LAUNCH_SOFTMAX(T, 8, WARPS_PER_TB); LAUNCH_SOFTMAX(T, 8, kWarpsPerTb);
break; break;
case 16: case 16:
LAUNCH_SOFTMAX(T, 16, WARPS_PER_TB); LAUNCH_SOFTMAX(T, 16, kWarpsPerTb);
break; break;
case 32: case 32:
LAUNCH_SOFTMAX(T, 32, WARPS_PER_TB); LAUNCH_SOFTMAX(T, 32, kWarpsPerTb);
break; break;
case 64: case 64:
LAUNCH_SOFTMAX(T, 64, WARPS_PER_TB); LAUNCH_SOFTMAX(T, 64, kWarpsPerTb);
break; break;
case 128: case 128:
LAUNCH_SOFTMAX(T, 128, WARPS_PER_TB); LAUNCH_SOFTMAX(T, 128, kWarpsPerTb);
break; break;
case 256: case 256:
LAUNCH_SOFTMAX(T, 256, WARPS_PER_TB); LAUNCH_SOFTMAX(T, 256, kWarpsPerTb);
break; break;
default: { default: {
TORCH_CHECK(softmax_workspace != nullptr, "softmax_workspace must be provided for num_experts that are not a power of 2."); CHECK(softmax_workspace != nullptr)
static constexpr int TPB = 256; << "softmax_workspace must be provided for num_experts that are "
moe_softmax<T, TPB><<<num_tokens, TPB, 0, stream>>>(gating_output, "not a power of 2.";
nullptr, static constexpr int kTpb = 256;
softmax_workspace, moe_softmax<T, kTpb><<<num_tokens, kTpb, 0, stream>>>(gating_output,
num_experts, nullptr,
moe_softcapping, softmax_workspace,
correction_bias); num_experts,
moe_softcapping,
correction_bias);
if (topk == 1) { if (topk == 1) {
// Note: As an optimization for better performance, // Note: As an optimization for better performance,
// the softmax_workspace is overwritten in-place by both moeTopK and // the softmax_workspace is overwritten in-place by both moeTopK and
// moe_topk_fast. // moe_topk_fast.
moe_topK<TPB><<<num_tokens, TPB, 0, stream>>>(softmax_workspace, moe_topK<kTpb><<<num_tokens, kTpb, 0, stream>>>(softmax_workspace,
nullptr, nullptr,
topk_weights, topk_weights,
topk_indices, topk_indices,
num_experts, num_experts,
topk, topk,
0, 0,
num_experts, num_experts,
renormalize); renormalize);
} else { } else {
moe_topk_fast<TPB><<<num_tokens, TPB, 0, stream>>>(softmax_workspace, moe_topk_fast<kTpb><<<num_tokens, kTpb, 0, stream>>>(softmax_workspace,
nullptr, nullptr,
topk_weights, topk_weights,
topk_indices, topk_indices,
num_experts, num_experts,
topk, topk,
0, 0,
num_experts, num_experts,
renormalize); renormalize);
} }
} }
} }
@@ -749,20 +763,29 @@ void topk_softmax(torch::Tensor& topk_weights, // [num_tokens, topk]
const double moe_softcapping, const double moe_softcapping,
const std::optional<torch::Tensor>& correction_bias) { const std::optional<torch::Tensor>& correction_bias) {
// Check data type // 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::Half ||
gating_output.scalar_type() == at::ScalarType::BFloat16, gating_output.scalar_type() == at::ScalarType::BFloat16)
"gating_output must be float32, float16, or bfloat16"); << "gating_output must be float32, float16, or bfloat16";
// Check dimensions // Check dimensions
TORCH_CHECK(gating_output.dim() == 2, "gating_output must be 2D tensor [num_tokens, num_experts]"); CHECK(gating_output.dim() == 2)
TORCH_CHECK(topk_weights.dim() == 2, "topk_weights must be 2D tensor [num_tokens, topk]"); << "gating_output must be 2D tensor [num_tokens, num_experts]";
TORCH_CHECK(topk_indices.dim() == 2, "topk_indices must be 2D tensor [num_tokens, topk]"); 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 shapes
TORCH_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"); 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_experts = static_cast<int>(gating_output.size(-1));
const int num_tokens = static_cast<int>(gating_output.size(0)); const int num_tokens = static_cast<int>(gating_output.size(0));
@@ -784,9 +807,12 @@ void topk_softmax(torch::Tensor& topk_weights, // [num_tokens, topk]
const float* bias_ptr = nullptr; const float* bias_ptr = nullptr;
if (correction_bias.has_value()) { if (correction_bias.has_value()) {
const torch::Tensor& bias_tensor = correction_bias.value(); const torch::Tensor& bias_tensor = correction_bias.value();
TORCH_CHECK(bias_tensor.dim() == 1, "correction_bias must be 1D tensor [num_experts]"); CHECK(bias_tensor.dim() == 1)
TORCH_CHECK(bias_tensor.size(0) == num_experts, "correction_bias size must match num_experts"); << "correction_bias must be 1D tensor [num_experts]";
TORCH_CHECK(bias_tensor.scalar_type() == at::ScalarType::Float, "correction_bias must be float32"); 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>(); bias_ptr = bias_tensor.data_ptr<float>();
} }
@@ -820,8 +846,8 @@ void topk_softmax(torch::Tensor& topk_weights, // [num_tokens, topk]
bias_ptr, bias_ptr,
stream); stream);
} else if (dtype == at::ScalarType::BFloat16) { } else if (dtype == at::ScalarType::BFloat16) {
topk_gating_softmax_kernel_launcher<__nv_bfloat16>( topk_gating_softmax_kernel_launcher<BFloat16Type>(
reinterpret_cast<const __nv_bfloat16*>( reinterpret_cast<const BFloat16Type*>(
gating_output.data_ptr<at::BFloat16>()), gating_output.data_ptr<at::BFloat16>()),
topk_weights.data_ptr<float>(), topk_weights.data_ptr<float>(),
topk_indices.data_ptr<int>(), topk_indices.data_ptr<int>(),
@@ -834,7 +860,7 @@ void topk_softmax(torch::Tensor& topk_weights, // [num_tokens, topk]
bias_ptr, bias_ptr,
stream); stream);
} else { } else {
TORCH_CHECK(false, "Unsupported gating_output dtype"); LOG(FATAL) << "Unsupported gating_output dtype: " << dtype;
} }
} }
} // namespace xllm::kernel::cuda } // namespace xllm::kernel::cuda

BIN
qwen3_6_scripts/prebuilt/corex-3.2.3-ivcore10/xllm_fused_qknorm_rope.so Executable file
View File

Binary file not shown.

231
qwen3_6_scripts/qwen3_5.py
View File

@@ -7,6 +7,7 @@ import hashlib
import os import os
import sys import sys
import time import time
print("[qwen3_5] module load START", file=sys.stderr, flush=True)
from typing import (Any, Dict, Iterable, List, Literal, Mapping, Optional, from typing import (Any, Dict, Iterable, List, Literal, Mapping, Optional,
Tuple, TypedDict, Union) Tuple, TypedDict, Union)
@@ -166,6 +167,41 @@ except ImportError:
except ImportError: except ImportError:
_xllm_moe = None _xllm_moe = None
# --- xllm prebuilt kernel loading (PRD build) ---
def _load_xllm_prebuilt(name):
"""Load a prebuilt xllm .so from corex-3.2.3-ivcore10 directory."""
import importlib.util as _ilu
_search = [
f"/workspace/qwen3_6_scripts/prebuilt/corex-3.2.3-ivcore10/{name}.so",
os.path.join(os.path.dirname(__file__), "prebuilt",
"corex-3.2.3-ivcore10", f"{name}.so"),
# When patch_ops.sh copies this file into vllm package, __file__
# points to vllm/model_executor/models/ — look back up to workspace
f"/home/dylan/0814/project_6/qwen3_6_scripts/prebuilt/corex-3.2.3-ivcore10/{name}.so",
]
for _p in _search:
if os.path.isfile(_p):
print(f"[xllm] loading {name} from {_p} ...", file=sys.stderr, flush=True)
try:
_spec = _ilu.spec_from_file_location(name, _p)
_mod = _ilu.module_from_spec(_spec)
_spec.loader.exec_module(_mod)
print(f"[xllm] {name} OK: {[x for x in dir(_mod) if not x.startswith('_')]}", file=sys.stderr, flush=True)
return _mod
except Exception as _e:
print(f"[xllm] {name} FAILED: {_e}", file=sys.stderr, flush=True)
return None
print(f"[xllm] {name} not found", file=sys.stderr, flush=True)
return None
print("[xllm] loading prebuilt kernels ...", file=sys.stderr, flush=True)
_xllm_norm = _load_xllm_prebuilt("xllm_norm")
_xllm_rope = _load_xllm_prebuilt("xllm_rope")
_xllm_activation = _load_xllm_prebuilt("xllm_activation")
_xllm_cache = _load_xllm_prebuilt("xllm_cache")
_xllm_fused_qknorm_rope = _load_xllm_prebuilt("xllm_fused_qknorm_rope")
print("[xllm] prebuilt kernel loading done", file=sys.stderr, flush=True)
# ix_moe_bridge: direct GEMV (4.1x faster than F.linear for decode M=1) # ix_moe_bridge: direct GEMV (4.1x faster than F.linear for decode M=1)
# Benchmark: F.linear 133us vs br.linear 32us on BI-V100 # Benchmark: F.linear 133us vs br.linear 32us on BI-V100
try: try:
@@ -189,7 +225,7 @@ _HAS_BRIDGE_LINEAR = (
_ix_moe_bridge is not None _ix_moe_bridge is not None
and hasattr(_ix_moe_bridge, 'linear')) and hasattr(_ix_moe_bridge, 'linear'))
if _HAS_BRIDGE_LINEAR: if _HAS_BRIDGE_LINEAR:
logger.info("ix_moe_bridge.linear ENABLED — 4.1x GEMV speedup for decode") print("[xllm] ix_moe_bridge.linear ENABLED", file=sys.stderr, flush=True)
def _fast_linear(x: torch.Tensor, weight: torch.Tensor, def _fast_linear(x: torch.Tensor, weight: torch.Tensor,
@@ -269,6 +305,141 @@ if _USE_XLLM_MOE:
logger.info("xllm_moe ENABLED — fused_topk + compute_index + combine_result") logger.info("xllm_moe ENABLED — fused_topk + compute_index + combine_result")
_USE_FUSED_MOE_ACTIVATION = env_bool("BI100_MOE_FUSED_ACTIVATION", True) _USE_FUSED_MOE_ACTIVATION = env_bool("BI100_MOE_FUSED_ACTIVATION", True)
# --- xllm CUDA kernel flags ---
_USE_XLLM_NORM = (
_xllm_norm is not None
and env_bool("BI100_XLLM_NORM", True))
_USE_XLLM_ROPE = (
_xllm_rope is not None
and env_bool("BI100_XLLM_ROPE", True))
_USE_XLLM_ACTIVATION = (
_xllm_activation is not None
and env_bool("BI100_XLLM_ACTIVATION", True))
_USE_XLLM_CACHE = (
_xllm_cache is not None
and env_bool("BI100_XLLM_CACHE", True))
_USE_XLLM_FUSED_QKNORM_ROPE = (
_xllm_fused_qknorm_rope is not None
and env_bool("BI100_XLLM_FUSED_QKNORM_ROPE", True))
if _USE_XLLM_NORM:
logger.info("xllm_norm ENABLED — fused RMSNorm CUDA kernel")
if _USE_XLLM_ROPE:
logger.info("xllm_rope ENABLED — fused RoPE CUDA kernel")
if _USE_XLLM_ACTIVATION:
logger.info("xllm_activation ENABLED — fused SiLU-and-Mul CUDA kernel")
if _USE_XLLM_CACHE:
logger.info("xllm_cache ENABLED — fused reshape_paged_cache CUDA kernel")
if _USE_XLLM_FUSED_QKNORM_ROPE:
logger.info("xllm_fused_qknorm_rope ENABLED — fused QK-Norm+RoPE (saves 128 launches/fwd)")
# ---------------------------------------------------------------------------
# xllm kernel monkey-patches: replace PyTorch fallback → fused CUDA kernels
# ---------------------------------------------------------------------------
print("[xllm] applying monkey-patches ...", file=sys.stderr, flush=True)
# --- 1. RMSNorm: xllm_norm replaces GemmaRMSNorm.forward_cuda ---
# GemmaRMSNorm uses x * (1 + weight) while xllm kernel uses x * weight.
# We pre-compute (1 + weight) and cache it so the kernel sees the correct
# effective weight. The cached tensor is lazily materialised on first call
# and invalidated if shape/device/dtype change (should never happen after
# model init).
if _USE_XLLM_NORM:
from vllm.model_executor.layers.layernorm import GemmaRMSNorm as _GemmaRMSNorm
def _xllm_rms_norm_forward_cuda(self, x, residual=None):
"""Drop-in for GemmaRMSNorm.forward_cuda using xllm_norm CUDA kernel."""
# Lazily compute and cache effective_weight = 1 + weight
_ew = getattr(self, '_xllm_eff_weight', None)
if (_ew is None
or _ew.shape != self.weight.shape
or _ew.device != self.weight.device
or _ew.dtype != self.weight.dtype):
_ew = (1.0 + self.weight.data.float()).to(self.weight.dtype)
self._xllm_eff_weight = _ew
if residual is not None:
# fused_add_rms_norm: input = RMSNorm(input + residual),
# residual = input + residual (before norm)
# The kernel modifies input and residual IN-PLACE.
_xllm_norm.fused_add_rms_norm(
x, residual, _ew, self.variance_epsilon)
return x, residual
else:
out = torch.empty_like(x)
_xllm_norm.rms_norm(out, x, _ew, self.variance_epsilon)
return out
_GemmaRMSNorm.forward_cuda = _xllm_rms_norm_forward_cuda
logger.info("xllm_norm PATCHED — GemmaRMSNorm.forward_cuda → xllm CUDA kernel")
# --- 2. SiluAndMul: xllm_activation replaces SiluAndMul.forward_cuda ---
if _USE_XLLM_ACTIVATION:
def _xllm_silu_and_mul_forward_cuda(self, x):
"""Drop-in for SiluAndMul.forward_cuda using xllm_activation kernel."""
d = x.shape[-1] // 2
output_shape = x.shape[:-1] + (d,)
# Reuse cached output tensor for stable shapes (decode)
_cache_key = (output_shape, x.dtype, x.device)
_cached = getattr(self, '_out_cache', {}).get(_cache_key)
if _cached is not None and _cached.shape == output_shape:
out = _cached
else:
out = torch.empty(output_shape, dtype=x.dtype, device=x.device)
if not hasattr(self, '_out_cache'):
self._out_cache = {}
self._out_cache[_cache_key] = out
_xllm_activation.silu_and_mul(out, x)
return out
SiluAndMul.forward_cuda = _xllm_silu_and_mul_forward_cuda
logger.info("xllm_activation PATCHED — SiluAndMul.forward_cuda → xllm CUDA kernel")
# --- 3. Cache ops: xllm_cache (reshape_paged_cache / block_copy) ---
# Replace ixformer vllm_cache_ops_reshape_and_cache with xllm_cache.
# xllm_cache signature: reshape_paged_cache(slot_ids, keys, values, kc, vc)
# vllm calls: ops.reshape_and_cache(key, value, kc, vc, slot_mapping, ...)
# Difference: arg order, slot_ids must be int32.
if _USE_XLLM_CACHE:
import vllm._custom_ops as _vllm_ops
_orig_reshape_and_cache = _vllm_ops.reshape_and_cache
def _xllm_reshape_and_cache(key, value, key_cache, value_cache,
slot_mapping, kv_cache_dtype="auto",
k_scale=1.0, v_scale=1.0):
slot_ids = slot_mapping.flatten().to(torch.int32)
_xllm_cache.reshape_paged_cache(slot_ids, key, value,
key_cache, value_cache)
_vllm_ops.reshape_and_cache = _xllm_reshape_and_cache
logger.info("xllm_cache PATCHED — reshape_and_cache → xllm CUDA kernel")
# --- 4. RoPE: xllm_rope ---
# Qwen3.5 uses interleaved multi-axis RoPE (MRoPE) with partial rotary
# factor 0.25. The xllm_rope kernel accepts (positions, query, key,
# cos_sin_cache, is_neox) and modifies q/k in-place, but it applies RoPE
# to the FULL head dim. Qwen3.5's forward only rotates the first 25% of
# each head (rotary_dim = 64 out of 256), then concatenates the unrotated
# tail. The multi-axis interleaving of cos/sin across T/H/W axes also
# requires Python-level assembly before calling any kernel.
#
# A safe replacement would need to:
# 1. Assemble the interleaved cos_sin for the partial dim.
# 2. Call xllm_rope on just the rotary_dim slice of q and k.
# 3. Skip the concat step since the kernel modifies in-place.
#
# This is doable but requires careful per-position indexing that differs
# from the standard "positions → cos_sin_cache[positions]" pattern.
# We mark this as TODO and leave the current _apply_rotary_emb path.
if _USE_XLLM_ROPE:
logger.info("xllm_rope LOADED but NOT PATCHED — Qwen3.5 interleaved MRoPE "
"requires adapter; using PyTorch _apply_rotary_emb fallback")
# ix_fused_moe: full 7-step fused MoE pipeline via ixformer C++ API # ix_fused_moe: full 7-step fused MoE pipeline via ixformer C++ API
# Source: xllm/core/layers/ilu/fused_moe.cpp → ix_moe_bridge.so # Source: xllm/core/layers/ilu/fused_moe.cpp → ix_moe_bridge.so
try: try:
@@ -1578,10 +1749,6 @@ class Qwen3_5FullAttention(nn.Module):
q = qg[:, :, :self.head_dim].reshape(total_tokens, -1) q = qg[:, :, :self.head_dim].reshape(total_tokens, -1)
gate = qg[:, :, self.head_dim:].reshape(total_tokens, -1) gate = qg[:, :, self.head_dim:].reshape(total_tokens, -1)
q = self.q_norm.forward_cuda(
q.view(total_tokens, self.local_num_heads, self.head_dim)
.contiguous()).view(total_tokens, -1)
# Select the one rank-local KV head before k_norm and RoPE. # Select the one rank-local KV head before k_norm and RoPE.
if self.q_per_kv_global is not None: if self.q_per_kv_global is not None:
tp_rank = get_tensor_model_parallel_rank() tp_rank = get_tensor_model_parallel_rank()
@@ -1592,10 +1759,53 @@ class Qwen3_5FullAttention(nn.Module):
v = (v.view(total_tokens, self.proj_kv_heads, self.head_dim) v = (v.view(total_tokens, self.proj_kv_heads, self.head_dim)
[:, kv_idx, :].contiguous()) [:, kv_idx, :].contiguous())
k = self.k_norm.forward_cuda( # --- Fused QK-Norm + RoPE path (saves 4 kernel launches per layer) ---
k.view(total_tokens, self.local_num_kv_heads, self.head_dim) # Only for 1D positions (decode / text-only prefill).
.contiguous()).view(total_tokens, -1) # 2D MRoPE positions (vision prefill) fall back to separate ops.
q, k = self.rotary_emb(positions, q, k) if (_USE_XLLM_FUSED_QKNORM_ROPE
and positions.ndim == 1
and q.is_contiguous() and k.is_contiguous()):
# Pack Q, K, V into contiguous [T, (Hq+Hk+Hv)*D] for fused kernel
v_flat = v.view(total_tokens, -1)
qkv = torch.cat([q, k.view(total_tokens, -1), v_flat], dim=-1)
# GemmaRMSNorm weight convention: kernel uses x*w, Gemma uses x*(1+w)
_q_ew = getattr(self, '_fused_q_ew', None)
if _q_ew is None:
_q_ew = (1.0 + self.q_norm.weight.data.float()).to(
self.q_norm.weight.dtype)
self._fused_q_ew = _q_ew
_k_ew = getattr(self, '_fused_k_ew', None)
if _k_ew is None:
_k_ew = (1.0 + self.k_norm.weight.data.float()).to(
self.k_norm.weight.dtype)
self._fused_k_ew = _k_ew
_xllm_fused_qknorm_rope.fused_qk_norm_rope(
qkv,
self.local_num_heads,
self.local_num_kv_heads,
self.local_num_kv_heads,
self.head_dim,
self.rms_norm_eps,
_q_ew,
_k_ew,
self.rotary_emb.cos_sin_cache,
True, # interleaved (Qwen3.5 uses interleaved RoPE)
positions.to(torch.int64))
# Unpack
q_dim = self.local_num_heads * self.head_dim
k_dim = self.local_num_kv_heads * self.head_dim
q = qkv[:, :q_dim]
k = qkv[:, q_dim:q_dim + k_dim]
# v is untouched by fused kernel, keep original
else:
# Fallback: separate q_norm, k_norm, rotary_emb
q = self.q_norm.forward_cuda(
q.view(total_tokens, self.local_num_heads, self.head_dim)
.contiguous()).view(total_tokens, -1)
k = self.k_norm.forward_cuda(
k.view(total_tokens, self.local_num_kv_heads, self.head_dim)
.contiguous()).view(total_tokens, -1)
q, k = self.rotary_emb(positions, q, k)
with bi100_timer("full_attn.attention"): with bi100_timer("full_attn.attention"):
with bi100_timer(f"L{self.layer_idx}.full_attn"): with bi100_timer(f"L{self.layer_idx}.full_attn"):
@@ -2889,4 +3099,5 @@ class Qwen3_5MoeForCausalLM(Qwen3_5ForCausalLM):
weight_loader(param, loaded_weight) weight_loader(param, loaded_weight)
_bi100_model_trace( _bi100_model_trace(
f"MoE load_weights complete items={loaded_count} " f"MoE load_weights complete items={loaded_count} "
f"vision_items={vision_loaded_count}") f"vision_items={vision_loaded_count}")
print("[qwen3_5] module load COMPLETE", file=sys.stderr, flush=True)

13
vllm/model_executor/models/mamba_cache.py
View File

@@ -10,11 +10,11 @@ class MambaCacheManager:
def __init__(self, dtype, num_mamba_layers, max_batch_size, def __init__(self, dtype, num_mamba_layers, max_batch_size,
conv_state_shape, temporal_state_shape): conv_state_shape, temporal_state_shape):
conv_state = torch.empty(size=(num_mamba_layers, max_batch_size) + conv_state = torch.zeros(size=(num_mamba_layers, max_batch_size) +
conv_state_shape, conv_state_shape,
dtype=dtype, dtype=dtype,
device="cuda") device="cuda")
temporal_state = torch.empty(size=(num_mamba_layers, max_batch_size) + temporal_state = torch.zeros(size=(num_mamba_layers, max_batch_size) +
temporal_state_shape, temporal_state_shape,
dtype=dtype, dtype=dtype,
device="cuda") device="cuda")
@@ -101,6 +101,8 @@ class MambaCacheManager:
self._move_out_if_already_occupied( self._move_out_if_already_occupied(
index=destination_index, index=destination_index,
all_occupied_indices=all_occupied_indices) all_occupied_indices=all_occupied_indices)
for cache_t in self.mamba_cache:
cache_t[:, destination_index].zero_()
self.mamba_cache_indices_mapping[cur_rid] = { self.mamba_cache_indices_mapping[cur_rid] = {
seq_id: destination_index seq_id: destination_index
} }
@@ -206,7 +208,10 @@ class MambaCacheManager:
finished_seq_groups_req_ids: List[str]): finished_seq_groups_req_ids: List[str]):
for req_id in finished_seq_groups_req_ids: for req_id in finished_seq_groups_req_ids:
if req_id in self.mamba_cache_indices_mapping: if req_id in self.mamba_cache_indices_mapping:
self.mamba_cache_indices_mapping.pop(req_id) seq_mapping = self.mamba_cache_indices_mapping.pop(req_id)
for cache_idx in seq_mapping.values():
for cache_t in self.mamba_cache:
cache_t[:, cache_idx].zero_()
def _first_free_index_in_mamba_cache( def _first_free_index_in_mamba_cache(
self, indices_range: Optional[List[int]] = None) -> int: self, indices_range: Optional[List[int]] = None) -> int:
@@ -219,4 +224,4 @@ class MambaCacheManager:
if i not in all_occupied_indices: if i not in all_occupied_indices:
return i return i
raise Exception("Couldn't find a free spot in the mamba cache! This" raise Exception("Couldn't find a free spot in the mamba cache! This"
"should never happen") "should never happen")

3313
vllm/model_executor/models/qwen3_5.py
View File

File diff suppressed because it is too large Load Diff

26
vllm/sequence.py
View File

@@ -119,6 +119,7 @@ class RequestMetrics:
scheduler_time: Optional[float] = None scheduler_time: Optional[float] = None
model_forward_time: Optional[float] = None model_forward_time: Optional[float] = None
model_execute_time: Optional[float] = None model_execute_time: Optional[float] = None
num_cached_tokens: Optional[int] = None
class SequenceDataDelta( class SequenceDataDelta(
@@ -527,6 +528,11 @@ class Sequence:
self._last_output_token_ids_offset = output_len self._last_output_token_ids_offset = output_len
# Return new tokens # Return new tokens
if num_new_tokens == 0:
# During chunked prefill steps with no output yet, num_new_tokens=0.
# Python's [-0:] == [0:] returns the ENTIRE list — guard against this.
return []
if num_new_tokens == 1: if num_new_tokens == 1:
# Optimization for single decode token case # Optimization for single decode token case
# (which is what we have most of the time) # (which is what we have most of the time)
@@ -935,6 +941,12 @@ class SequenceGroupMetadataDelta(
computed_block_nums: Optional[List[int]] = None computed_block_nums: Optional[List[int]] = None
state: Optional[SequenceGroupState] = msgspec.field( state: Optional[SequenceGroupState] = msgspec.field(
default_factory=lambda: SequenceGroupState()) default_factory=lambda: SequenceGroupState())
# BI100 hybrid prefix-cache actions. Fields are appended for msgspec wire
# compatibility with the pre-existing array-like structure.
gdn_restore_key: Optional[Tuple[int, bytes]] = None
gdn_capture_points: Optional[List[Tuple[int, Tuple[int, bytes]]]] = None
gdn_evict_keys: Optional[List[Tuple[int, bytes]]] = None
gdn_segment_offsets: Optional[List[int]] = None
class SequenceGroupMetadata( class SequenceGroupMetadata(
@@ -1000,6 +1012,12 @@ class SequenceGroupMetadata(
# Zero means speculative decoding is disabled for some reasons. # Zero means speculative decoding is disabled for some reasons.
# TODO: We should maintain this states out of the sequence group. # TODO: We should maintain this states out of the sequence group.
num_speculative_tokens: Optional[int] = None num_speculative_tokens: Optional[int] = None
# BI100 hybrid prefix-cache actions. These are internal scheduler-to-worker
# metadata and never surface through the OpenAI API.
gdn_restore_key: Optional[Tuple[int, bytes]] = None
gdn_capture_points: Optional[List[Tuple[int, Tuple[int, bytes]]]] = None
gdn_evict_keys: Optional[List[Tuple[int, bytes]]] = None
gdn_segment_offsets: Optional[List[int]] = None
def __post_init__(self): def __post_init__(self):
if self.seq_data is not None and self.token_chunk_size is None: if self.seq_data is not None and self.token_chunk_size is None:
@@ -1046,6 +1064,14 @@ class SequenceGroupMetadata(
self.token_chunk_size = sequence_group_metadata_delta.token_chunk_size self.token_chunk_size = sequence_group_metadata_delta.token_chunk_size
self.do_sample = sequence_group_metadata_delta.do_sample self.do_sample = sequence_group_metadata_delta.do_sample
self.is_prompt = sequence_group_metadata_delta.is_prompt self.is_prompt = sequence_group_metadata_delta.is_prompt
self.computed_block_nums = (
sequence_group_metadata_delta.computed_block_nums)
self.gdn_restore_key = sequence_group_metadata_delta.gdn_restore_key
self.gdn_capture_points = (
sequence_group_metadata_delta.gdn_capture_points)
self.gdn_evict_keys = sequence_group_metadata_delta.gdn_evict_keys
self.gdn_segment_offsets = (
sequence_group_metadata_delta.gdn_segment_offsets)
def finish_step(self) -> None: def finish_step(self) -> None:
assert self.state is not None assert self.state is not None

7
vllm/worker/worker.py
View File

@@ -216,6 +216,13 @@ class Worker(LocalOrDistributedWorkerBase):
""" """
# Profile the memory usage of the model and get the maximum number of # Profile the memory usage of the model and get the maximum number of
# cache blocks that can be allocated with the remaining free memory. # cache blocks that can be allocated with the remaining free memory.
# PRD: skip profile_run when num_gpu_blocks_override is set
_ovr = getattr(self.cache_config, 'num_gpu_blocks_override', None)
if _ovr is not None and _ovr > 0:
logger.info("Skipping profile_run -- num_gpu_blocks_override=%d", _ovr)
_cbs = self.get_cache_block_size_bytes()
_cpu = self.cache_config.swap_space_bytes // _cbs if _cbs > 0 else 256
return int(_ovr), int(_cpu)
torch.cuda.empty_cache() torch.cuda.empty_cache()
# Execute a forward pass with dummy inputs to profile the memory usage # Execute a forward pass with dummy inputs to profile the memory usage