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

feat(CRITICAL): ix_bridge call chain + upstream xllm/ds_vllm sync

Browse Source 
3 changes that close the MoE performance gap:

1. _custom_ops.py: Add ix_bridge as Priority 0 for topk_softmax
   - Before: tries our .cu kernel (fails) → PyTorch fallback (1-3 TPS)
   - After: tries ix_bridge → ixformer::infer::topk_softmax() → FAST
   - Call chain: _custom_ops.topk_softmax() → ix_bridge.topk_softmax()
     → ix_moe_bridge.so → ixformer::infer::topk_softmax()

2. Dockerfile: Add ix_moe_bridge.cpp precompile step
   - This was the missing link: code existed but was never compiled
   - Uses torch.utils.cpp_extension.load() to link against libixformer.so

3. upstream_ref sync from GitHub (cloned, not rewritten):
   - xLLM-AI/xllm: ILU kernels + CUDA MoE + GDN fp32 state mgmt
   - Deep-Spark/vllm: latest MoE kernel sources
This commit is contained in:
project6-dev
2026-08-11 01:27:33 +00:00
parent 26e6cb4019
commit 56146f8130
16 changed files with 5755 additions and 1 deletions
Download Patch File Download Diff File Expand all files Collapse all files

8
Dockerfile
View File

@@ -26,7 +26,13 @@ RUN chmod +x /workspace/qwen3_6_scripts/patch_ops.sh && \
bash /workspace/qwen3_6_scripts/patch_ops.sh 2>&1 | tee /workspace/patch_ops.log ; \
echo "[Dockerfile] patch_ops exit code: $?"
# Step 5: Precompile GDN kernel (needs vllm in path, so after patch_ops)
# Step 5: Precompile ix_moe_bridge.cpp → links to ixformer::infer::topk_softmax()
# This is the C++ pybind bridge that makes ixformer SDK callable from Python.
# Source: upstream_ref/xllm/xllm/core/kernels/ilu/fused_moe.cpp call pattern
RUN python3 /workspace/ex_engine/precompile_ix_bridge.py 2>&1 | tee -a /workspace/ex_build.log ; \
echo "[Dockerfile] ix_bridge precompile exit code: $?"
# Step 6: Precompile GDN kernel (needs vllm in path, so after patch_ops)
RUN python3 /workspace/qwen3_6_scripts/precompile_gdn.py \
/workspace/qwen3_6_scripts/flash_qla_sm70 2>&1 | tee -a /workspace/ex_build.log ; \
echo "[Dockerfile] gdn precompile exit code: $?"

107
ex_engine/precompile_ix_bridge.py Normal file
View File

@@ -0,0 +1,107 @@
#!/usr/bin/env python3
"""
Precompile ix_moe_bridge.cpp during Docker build.
This bridges Python ↔ ixformer::infer C++ API (topk_softmax, group_gemm, etc).
Source: upstream_ref/xllm/xllm/core/kernels/ilu/fused_moe.cpp call pattern
"""
import os
import sys
import glob
import logging
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger("precompile_ix_bridge")
def find_ixformer_libs():
"""Find libixformer.so and related libraries for linking."""
extra_ldflags = []
ixf_lib_dirs = set()
try:
import ixformer
ixf_dir = os.path.dirname(ixformer.__file__)
for so in glob.glob(os.path.join(ixf_dir, "*.so")):
if "cpython" not in so:
extra_ldflags.append(so)
ixf_lib_dirs.add(os.path.dirname(so))
for so in glob.glob(os.path.join(ixf_dir, "_ixformer_torch*.so")):
if so not in extra_ldflags:
extra_ldflags.append(so)
except ImportError:
logger.warning("ixformer not installed")
corex_lib = "/usr/local/corex/lib64"
if os.path.isdir(corex_lib):
for lib in ["libixformer.so", "libixattn.so", "libcublas.so"]:
p = os.path.join(corex_lib, lib)
if os.path.exists(p) and p not in extra_ldflags:
extra_ldflags.append(p)
ixf_lib_dirs.add(corex_lib)
for d in ixf_lib_dirs:
extra_ldflags.append(f"-Wl,-rpath,{d}")
return extra_ldflags
def main():
# Find the .cpp source
search = [
"/workspace/ex_engine/csrc/ix_moe_bridge.cpp",
os.path.join(os.path.dirname(__file__), "csrc", "ix_moe_bridge.cpp"),
]
cpp_path = None
for p in search:
if os.path.isfile(p):
cpp_path = p
break
if not cpp_path:
logger.error("ix_moe_bridge.cpp not found in: %s", search)
sys.exit(1)
logger.info("Compiling ix_moe_bridge from %s", cpp_path)
extra_ldflags = find_ixformer_libs()
logger.info("Link flags: %s", extra_ldflags)
if not extra_ldflags:
logger.error("No ixformer libraries found — cannot compile bridge")
sys.exit(1)
try:
from torch.utils.cpp_extension import load
mod = load(
name="ix_moe_bridge",
sources=[cpp_path],
extra_cflags=["-O2", "-std=c++17"],
extra_ldflags=extra_ldflags,
verbose=True,
)
fns = [x for x in dir(mod) if not x.startswith("_")]
logger.info("SUCCESS: ix_moe_bridge compiled with functions: %s", fns)
except Exception as e:
logger.error("FAILED to compile ix_moe_bridge: %s", e)
# Also try ix_full_bridge.cpp
full_path = cpp_path.replace("ix_moe_bridge", "ix_full_bridge")
if os.path.isfile(full_path):
logger.info("Trying ix_full_bridge.cpp instead...")
try:
mod = load(
name="ix_full_bridge",
sources=[full_path],
extra_cflags=["-O2", "-std=c++17"],
extra_ldflags=extra_ldflags,
verbose=True,
)
fns = [x for x in dir(mod) if not x.startswith("_")]
logger.info("SUCCESS: ix_full_bridge compiled with functions: %s", fns)
except Exception as e2:
logger.error("FAILED ix_full_bridge too: %s", e2)
sys.exit(1)
else:
sys.exit(1)
if __name__ == "__main__":
main()

50
qwen3_6_scripts/_custom_ops.py
View File

@@ -974,9 +974,44 @@ def invoke_fused_moe_kernel(
_moe_topk_ext = None
_moe_topk_init_done = False
_ix_bridge_mod = None
_ix_bridge_init_done = False
def _init_ix_bridge():
"""Try to load ix_bridge which calls ixformer::infer::topk_softmax() via C++ pybind."""
global _ix_bridge_mod, _ix_bridge_init_done
_ix_bridge_init_done = True
try:
from ex_engine.python.ix_bridge import is_available, topk_softmax as _ix_ts
if is_available():
_ix_bridge_mod = True
logger.info("topk_softmax: ix_bridge → ixformer::infer::topk_softmax() LOADED")
return
except Exception as e:
logger.info("topk_softmax: ix_bridge unavailable (%s)", e)
# Also try direct import from workspace
try:
import sys
for p in ['/workspace/ex_engine/python', '/workspace/ex_engine',
'/usr/local/corex/lib/python3/dist-packages/ex_engine/python']:
if p not in sys.path:
sys.path.insert(0, p)
from ix_bridge import is_available, topk_softmax as _ix_ts
if is_available():
_ix_bridge_mod = True
logger.info("topk_softmax: ix_bridge (direct) → ixformer::infer LOADED")
return
except Exception as e:
logger.info("topk_softmax: ix_bridge direct import failed (%s)", e)
def _init_moe_topk():
global _moe_topk_ext, _moe_topk_init_done
_moe_topk_init_done = True
# 0. Try ix_bridge first (calls ixformer C++ SDK directly)
_init_ix_bridge()
if _ix_bridge_mod:
return # ix_bridge loaded, no need for CUDA kernel
# 1. Try import precompiled module (torch cache from Docker build)
try:
import moe_topk_softmax_v3 as ext
@@ -1038,6 +1073,21 @@ def topk_softmax(topk_weights: torch.Tensor, topk_ids: torch.Tensor,
if not _moe_topk_init_done:
_init_moe_topk()
# Priority 0: ix_bridge → ixformer::infer::topk_softmax() (fastest, uses SDK)
if _ix_bridge_mod:
try:
from ex_engine.python.ix_bridge import topk_softmax as _ix_topk
gating = gating_output if isinstance(gating_output, torch.Tensor) else gating_output
topk_k = topk_weights.shape[1]
weights, ids = _ix_topk(gating, topk_k, renormalize=False)
topk_weights.copy_(weights.to(topk_weights.dtype))
topk_ids.copy_(ids.to(topk_ids.dtype))
# token_expert_indicies not produced by ix_bridge, fill with topk_ids
token_expert_indicies.copy_(ids.to(token_expert_indicies.dtype))
return
except Exception as e:
logger.warning("topk_softmax ix_bridge failed (%s), trying CUDA kernel", e)
# Priority 1: Our CUDA kernel (fused warp-shuffle, ~5x faster than PyTorch)
if _moe_topk_ext is not None:
try:

257
upstream_ref/ds_vllm_latest/csrc/moe/moeTopKFuncs.cuh Normal file
View File

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

833
upstream_ref/ds_vllm_latest/csrc/moe/moe_align_sum_kernels.cu Normal file
View File

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

87
upstream_ref/ds_vllm_latest/csrc/moe/moe_ops.h Normal file
View File

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

874
upstream_ref/ds_vllm_latest/csrc/moe/topk_softmax_kernels.cu Normal file
View File

@@ -0,0 +1,874 @@
/*
* Adapted from https://github.com/NVIDIA/TensorRT-LLM/blob/v0.7.1/cpp/tensorrt_llm/kernels/mixtureOfExperts/moe_kernels.cu
* Copyright (c) 2024, The vLLM team.
* SPDX-FileCopyrightText: Copyright (c) 1993-2023 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
* SPDX-License-Identifier: Apache-2.0
*
* Licensed under the Apache License, Version 2.0 (the "License");
* you may not use this file except in compliance with the License.
* You may obtain a copy of the License at
*
* http://www.apache.org/licenses/LICENSE-2.0
*
* Unless required by applicable law or agreed to in writing, software
* distributed under the License is distributed on an "AS IS" BASIS,
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
* See the License for the specific language governing permissions and
* limitations under the License.
*/
#include <type_traits>
#include <cuda_runtime.h>
#include <torch/csrc/stable/accelerator.h>
#include <torch/csrc/stable/tensor.h>
#include <torch/headeronly/core/ScalarType.h>
#include <torch/headeronly/util/Exception.h>
#include "../../cuda_compat.h"
#include "../../cub_helpers.h"
#include "libtorch_stable/torch_utils.h"
#ifndef USE_ROCM
#include <cuda_bf16.h>
#include <cuda_fp16.h>
#else
#include <hip/hip_bf16.h>
#include <hip/hip_fp16.h>
typedef __hip_bfloat16 __nv_bfloat16;
typedef __hip_bfloat162 __nv_bfloat162;
#endif
#define MAX(a, b) ((a) > (b) ? (a) : (b))
#define MIN(a, b) ((a) < (b) ? (a) : (b))
namespace vllm {
namespace moe {
/// Aligned array type
template <
typename T,
/// Number of elements in the array
int N,
/// Alignment requirement in bytes
int Alignment = sizeof(T) * N
>
struct alignas(Alignment) AlignedArray {
T data[N];
};
template <typename T>
__device__ __forceinline__ float toFloat(T value) {
if constexpr (std::is_same_v<T, float>) {
return value;
} else if constexpr (std::is_same_v<T, __nv_bfloat16>) {
return __bfloat162float(value);
} else if constexpr (std::is_same_v<T, __half>) {
return __half2float(value);
}
}
// Scoring function enums
enum ScoringFunc {
SCORING_SOFTMAX = 0, // apply softmax
SCORING_SIGMOID = 1 // apply sigmoid
};
// ====================== Softmax things ===============================
// We have our own implementation of softmax here so we can support transposing the output
// in the softmax kernel when we extend this module to support expert-choice routing.
template <int TPB, typename InputType>
__launch_bounds__(TPB) __global__
void moeSoftmax(const InputType* input, const bool* finished, float* output, const int num_cols)
{
using BlockReduce = cub::BlockReduce<float, TPB>;
__shared__ typename BlockReduce::TempStorage tmpStorage;
__shared__ float normalizing_factor;
__shared__ float float_max;
const int thread_row_offset = blockIdx.x * num_cols;
float threadData(-FLT_MAX);
// Don't touch finished rows.
if ((finished != nullptr) && finished[blockIdx.x])
{
return;
}
for (int ii = threadIdx.x; ii < num_cols; ii += TPB)
{
const int idx = thread_row_offset + ii;
const float val = toFloat(input[idx]);
threadData = max(val, threadData);
}
const float maxElem = BlockReduce(tmpStorage).Reduce(threadData, CubMaxOp());
if (threadIdx.x == 0)
{
float_max = maxElem;
}
__syncthreads();
threadData = 0;
for (int ii = threadIdx.x; ii < num_cols; ii += TPB)
{
const int idx = thread_row_offset + ii;
const float val = toFloat(input[idx]);
threadData += expf(val - float_max);
}
const auto Z = BlockReduce(tmpStorage).Reduce(threadData, CubAddOp());
if (threadIdx.x == 0)
{
normalizing_factor = 1.f / Z;
}
__syncthreads();
for (int ii = threadIdx.x; ii < num_cols; ii += TPB)
{
const int idx = thread_row_offset + ii;
const float val = toFloat(input[idx]);
float softmax_val = expf(val - float_max) * normalizing_factor;
// Clamp NaN/Inf to 0 to prevent duplicate expert IDs downstream.
if (isnan(softmax_val) || isinf(softmax_val)) softmax_val = 0.f;
output[idx] = softmax_val;
}
}
template <int TPB, typename InputType>
__launch_bounds__(TPB) __global__
void moeSigmoid(const InputType* input, const bool* finished, float* output, const int num_cols)
{
const int thread_row_offset = blockIdx.x * num_cols;
// Don't touch finished rows.
if ((finished != nullptr) && finished[blockIdx.x])
{
return;
}
for (int ii = threadIdx.x; ii < num_cols; ii += TPB)
{
const int idx = thread_row_offset + ii;
const float val = toFloat(input[idx]);
float sigmoid_val = 1.0f / (1.0f + __expf(-val));
// Clamp NaN/Inf to 0 to prevent duplicate expert IDs downstream.
if (isnan(sigmoid_val) || isinf(sigmoid_val)) sigmoid_val = 0.f;
output[idx] = sigmoid_val;
}
}
template <int TPB, typename IndType>
__launch_bounds__(TPB) __global__ void moeTopK(
const float* inputs_after_softmax,
const bool* finished,
float* output,
IndType* indices,
int* source_rows,
const int num_experts,
const int k,
const int start_expert,
const int end_expert,
const bool renormalize,
const float* bias)
{
using cub_kvp = cub::KeyValuePair<int, float>;
using BlockReduce = cub::BlockReduce<cub_kvp, TPB>;
__shared__ typename BlockReduce::TempStorage tmpStorage;
cub_kvp thread_kvp;
cub::ArgMax arg_max;
const int num_rows = gridDim.x;
const int block_row = blockIdx.x;
const bool row_is_active = finished ? !finished[block_row] : true;
const int thread_read_offset = blockIdx.x * num_experts;
float selected_sum = 0.f;
for (int k_idx = 0; k_idx < k; ++k_idx)
{
thread_kvp.key = 0;
thread_kvp.value = -1.f; // This is OK because inputs are probabilities
cub_kvp inp_kvp;
for (int expert = threadIdx.x; expert < num_experts; expert += TPB)
{
const int idx = thread_read_offset + expert;
inp_kvp.key = expert;
// Apply correction bias if provided
if (bias != nullptr) {
inp_kvp.value = inputs_after_softmax[idx] + bias[expert];
} else {
inp_kvp.value = inputs_after_softmax[idx];
}
for (int prior_k = 0; prior_k < k_idx; ++prior_k)
{
const int prior_winning_expert = indices[k * block_row + prior_k];
if (prior_winning_expert == expert)
{
inp_kvp = thread_kvp;
}
}
thread_kvp = arg_max(inp_kvp, thread_kvp);
}
const cub_kvp result_kvp = BlockReduce(tmpStorage).Reduce(thread_kvp, arg_max);
if (threadIdx.x == 0)
{
// Ignore experts the node isn't responsible for with expert parallelism
const int expert = result_kvp.key;
const bool node_uses_expert = expert >= start_expert && expert < end_expert;
const bool should_process_row = row_is_active && node_uses_expert;
const int idx = k * block_row + k_idx;
// Return the unbiased scores for output weights
output[idx] = inputs_after_softmax[thread_read_offset + expert];
indices[idx] = should_process_row ? (expert - start_expert) : num_experts;
assert(indices[idx] >= 0);
source_rows[idx] = k_idx * num_rows + block_row;
if (renormalize) {
selected_sum += inputs_after_softmax[thread_read_offset + expert];
}
}
__syncthreads();
}
// Renormalize the k weights for this row to sum to 1, if requested.
if (renormalize) {
if (threadIdx.x == 0) {
const float denom = selected_sum > 0.f ? selected_sum : 1.f;
for (int k_idx = 0; k_idx < k; ++k_idx) {
const int idx = k * block_row + k_idx;
output[idx] = output[idx] / denom;
}
}
}
}
// ====================== TopK softmax things ===============================
/*
A Top-K gating softmax written to exploit when the number of experts in the MoE layers
are a small power of 2. This allows us to cleanly share the rows among the threads in
a single warp and eliminate communication between warps (so no need to use shared mem).
It fuses the softmax, max and argmax into a single kernel.
Limitations:
1) This implementation is optimized for when the number of experts is a small power of 2.
Additionally it also supports when number of experts is multiple of 64 which is still
faster than the computing softmax and topK separately (only tested on CUDA yet).
2) This implementation assumes k is small, but will work for any k.
*/
template <int VPT, int NUM_EXPERTS, int WARPS_PER_CTA, int BYTES_PER_LDG, int WARP_SIZE_PARAM, typename IndType,
typename InputType = float, ScoringFunc SF>
__launch_bounds__(WARPS_PER_CTA* WARP_SIZE_PARAM) __global__
void topkGating(const InputType* input, const bool* finished, float* output, const int num_rows, IndType* indices,
int* source_rows, const int k, const int start_expert, const int end_expert, const bool renormalize,
const float* bias)
{
static_assert(std::is_same_v<InputType, float> || std::is_same_v<InputType, __nv_bfloat16> ||
std::is_same_v<InputType, __half>,
"InputType must be float, __nv_bfloat16, or __half");
// We begin by enforcing compile time assertions and setting up compile time constants.
static_assert(BYTES_PER_LDG == (BYTES_PER_LDG & -BYTES_PER_LDG), "BYTES_PER_LDG must be power of 2");
static_assert(BYTES_PER_LDG <= 16, "BYTES_PER_LDG must be leq 16");
// Number of bytes each thread pulls in per load
static constexpr int ELTS_PER_LDG = BYTES_PER_LDG / sizeof(InputType);
static constexpr int ELTS_PER_ROW = NUM_EXPERTS;
static constexpr int THREADS_PER_ROW = ELTS_PER_ROW / VPT;
static constexpr int LDG_PER_THREAD = VPT / ELTS_PER_LDG;
if constexpr (std::is_same_v<InputType, __nv_bfloat16> || std::is_same_v<InputType, __half>) {
static_assert(ELTS_PER_LDG == 1 || ELTS_PER_LDG % 2 == 0,
"ELTS_PER_LDG must be 1 or even for 16-bit conversion");
}
// Restrictions based on previous section.
static_assert(VPT % ELTS_PER_LDG == 0, "The elements per thread must be a multiple of the elements per ldg");
static_assert(WARP_SIZE_PARAM % THREADS_PER_ROW == 0, "The threads per row must cleanly divide the threads per warp");
static_assert(THREADS_PER_ROW == (THREADS_PER_ROW & -THREADS_PER_ROW), "THREADS_PER_ROW must be power of 2");
static_assert(THREADS_PER_ROW <= WARP_SIZE_PARAM, "THREADS_PER_ROW can be at most warp size");
// We have NUM_EXPERTS elements per row. We specialize for small #experts
static constexpr int ELTS_PER_WARP = WARP_SIZE_PARAM * VPT;
static constexpr int ROWS_PER_WARP = ELTS_PER_WARP / ELTS_PER_ROW;
static constexpr int ROWS_PER_CTA = WARPS_PER_CTA * ROWS_PER_WARP;
// Restrictions for previous section.
static_assert(ELTS_PER_WARP % ELTS_PER_ROW == 0, "The elts per row must cleanly divide the total elt per warp");
// ===================== From this point, we finally start computing run-time variables. ========================
// Compute CTA and warp rows. We pack multiple rows into a single warp, and a block contains WARPS_PER_CTA warps.
// This, each block processes a chunk of rows. We start by computing the start row for each block.
const int cta_base_row = blockIdx.x * ROWS_PER_CTA;
// Now, using the base row per thread block, we compute the base row per warp.
const int warp_base_row = cta_base_row + threadIdx.y * ROWS_PER_WARP;
// The threads in a warp are split into sub-groups that will work on a row.
// We compute row offset for each thread sub-group
const int thread_row_in_warp = threadIdx.x / THREADS_PER_ROW;
const int thread_row = warp_base_row + thread_row_in_warp;
// Threads with indices out of bounds should early exit here.
if (thread_row >= num_rows)
{
return;
}
const bool row_is_active = finished ? !finished[thread_row] : true;
// We finally start setting up the read pointers for each thread. First, each thread jumps to the start of the
// row it will read.
const InputType* thread_row_ptr = input + thread_row * ELTS_PER_ROW;
// Now, we compute the group each thread belong to in order to determine the first column to start loads.
const int thread_group_idx = threadIdx.x % THREADS_PER_ROW;
const int first_elt_read_by_thread = thread_group_idx * ELTS_PER_LDG;
const InputType* thread_read_ptr = thread_row_ptr + first_elt_read_by_thread;
// Finally, we pull in the data from global mem
float row_chunk[VPT];
// NOTE(zhuhaoran): dispatch different input types loading, BF16/FP16 convert to float
if constexpr (std::is_same_v<InputType, float>) {
using VecType = AlignedArray<float, ELTS_PER_LDG>;
VecType* row_chunk_vec_ptr = reinterpret_cast<VecType*>(&row_chunk);
const VecType* vec_thread_read_ptr = reinterpret_cast<const VecType*>(thread_read_ptr);
#pragma unroll
for (int ii = 0; ii < LDG_PER_THREAD; ++ii) {
row_chunk_vec_ptr[ii] = vec_thread_read_ptr[ii * THREADS_PER_ROW];
}
} else if constexpr (std::is_same_v<InputType, __nv_bfloat16>) {
if constexpr (ELTS_PER_LDG >= 2) {
using VecType = AlignedArray<__nv_bfloat16, ELTS_PER_LDG>;
float2* row_chunk_f2 = reinterpret_cast<float2*>(row_chunk);
const VecType* vec_thread_read_ptr = reinterpret_cast<const VecType*>(thread_read_ptr);
#pragma unroll
for (int ii = 0; ii < LDG_PER_THREAD; ++ii) {
VecType vec = vec_thread_read_ptr[ii * THREADS_PER_ROW];
int base_idx_f2 = ii * ELTS_PER_LDG / 2;
#pragma unroll
for (int jj = 0; jj < ELTS_PER_LDG / 2; ++jj) {
row_chunk_f2[base_idx_f2 + jj] = __bfloat1622float2(
*reinterpret_cast<const __nv_bfloat162*>(vec.data + jj * 2)
);
}
}
} else { // ELTS_PER_LDG == 1
#pragma unroll
for (int ii = 0; ii < LDG_PER_THREAD; ++ii) {
const __nv_bfloat16* scalar_ptr = thread_read_ptr + ii * THREADS_PER_ROW;
row_chunk[ii] = __bfloat162float(*scalar_ptr);
}
}
} else if constexpr (std::is_same_v<InputType, __half>) {
if constexpr (ELTS_PER_LDG >= 2) {
using VecType = AlignedArray<__half, ELTS_PER_LDG>;
float2* row_chunk_f2 = reinterpret_cast<float2*>(row_chunk);
const VecType* vec_thread_read_ptr = reinterpret_cast<const VecType*>(thread_read_ptr);
#pragma unroll
for (int ii = 0; ii < LDG_PER_THREAD; ++ii) {
VecType vec = vec_thread_read_ptr[ii * THREADS_PER_ROW];
int base_idx_f2 = ii * ELTS_PER_LDG / 2;
#pragma unroll
for (int jj = 0; jj < ELTS_PER_LDG / 2; ++jj) {
row_chunk_f2[base_idx_f2 + jj] = __half22float2(
*reinterpret_cast<const __half2*>(vec.data + jj * 2)
);
}
}
} else { // ELTS_PER_LDG == 1
#pragma unroll
for (int ii = 0; ii < LDG_PER_THREAD; ++ii) {
const __half* scalar_ptr = thread_read_ptr + ii * THREADS_PER_ROW;
row_chunk[ii] = __half2float(*scalar_ptr);
}
}
}
if constexpr (SF == SCORING_SOFTMAX) {
// First, we perform a max reduce within the thread.
float thread_max = row_chunk[0];
#pragma unroll
for (int ii = 1; ii < VPT; ++ii) {
thread_max = max(thread_max, row_chunk[ii]);
}
// Now, we find the max within the thread group and distribute among the threads. We use a butterfly reduce.
#pragma unroll
for (int mask = THREADS_PER_ROW / 2; mask > 0; mask /= 2)
{
thread_max = max(thread_max, VLLM_SHFL_XOR_SYNC_WIDTH(thread_max, mask, THREADS_PER_ROW));
}
// From this point, thread max in all the threads have the max within the row.
// Now, we subtract the max from each element in the thread and take the exp. We also compute the thread local sum.
float row_sum = 0;
#pragma unroll
for (int ii = 0; ii < VPT; ++ii)
{
row_chunk[ii] = expf(row_chunk[ii] - thread_max);
row_sum += row_chunk[ii];
}
// Now, we perform the sum reduce within each thread group. Similar to the max reduce, we use a bufferfly pattern.
#pragma unroll
for (int mask = THREADS_PER_ROW / 2; mask > 0; mask /= 2)
{
row_sum += VLLM_SHFL_XOR_SYNC_WIDTH(row_sum, mask, THREADS_PER_ROW);
}
// From this point, all threads have the max and the sum for their rows in the thread_max and thread_sum variables
// respectively. Finally, we can scale the rows for the softmax. Technically, for top-k gating we don't need to
// compute the entire softmax row. We can likely look at the maxes and only compute for the top-k values in the row.
// However, this kernel will likely not be a bottle neck and it seems better to closer match torch and find the
// argmax after computing the softmax.
const float reciprocal_row_sum = 1.f / row_sum;
#pragma unroll
for (int ii = 0; ii < VPT; ++ii)
{
row_chunk[ii] = row_chunk[ii] * reciprocal_row_sum;
}
} else if constexpr (SF == SCORING_SIGMOID) {
#pragma unroll
for (int ii = 0; ii < VPT; ++ii)
{
row_chunk[ii] = 1.0f / (1.0f + __expf(-row_chunk[ii]));
}
}
// Fix: clamp NaN/Inf values to 0 to prevent duplicate expert IDs.
// NaN gating (from degenerate hidden states in CUDA graph padding) causes
// softmax to produce all-NaN, which makes the argmax loop always pick
// expert 0 for every top-k slot, producing duplicate expert IDs that
// crash FlashInfer's three-step MoE sort.
// With 0s, the argmax uses index tie-breaking to pick [0,1,2,...,k-1].
#pragma unroll
for (int ii = 0; ii < VPT; ++ii) {
if (isnan(row_chunk[ii]) || isinf(row_chunk[ii])) {
row_chunk[ii] = 0.f;
}
}
static constexpr int COLS_PER_GROUP_LDG = ELTS_PER_LDG * THREADS_PER_ROW;
// If bias is not null, use biased value for selection
float row_chunk_for_choice[VPT];
// Apply correction bias
if (bias != nullptr) {
#pragma unroll
for (int ldg = 0; ldg < LDG_PER_THREAD; ++ldg) {
#pragma unroll
for (int ii = 0; ii < ELTS_PER_LDG; ++ii) {
const int expert = first_elt_read_by_thread + ldg * COLS_PER_GROUP_LDG + ii;
float bias_val = expert < NUM_EXPERTS ? bias[expert] : 0.0f;
row_chunk_for_choice[ldg * ELTS_PER_LDG + ii] = row_chunk[ldg * ELTS_PER_LDG + ii] + bias_val;
}
}
} else {
#pragma unroll
for (int ii = 0; ii < VPT; ++ii) {
row_chunk_for_choice[ii] = row_chunk[ii];
}
}
// Now, row_chunk contains the softmax / sigmoid of the row chunk. Now, I want to find the topk elements in each row, along
// with the max index.
int start_col = first_elt_read_by_thread;
float selected_sum = 0.f;
for (int k_idx = 0; k_idx < k; ++k_idx)
{
// First, each thread does the local argmax
float max_val_for_choice = row_chunk_for_choice[0];
float max_val = row_chunk[0];
int expert = start_col;
#pragma unroll
for (int ldg = 0, col = start_col; ldg < LDG_PER_THREAD; ++ldg, col += COLS_PER_GROUP_LDG)
{
#pragma unroll
for (int ii = 0; ii < ELTS_PER_LDG; ++ii)
{
float val_for_choice = row_chunk_for_choice[ldg * ELTS_PER_LDG + ii];
float val = row_chunk[ldg * ELTS_PER_LDG + ii];
// No check on the experts here since columns with the smallest index are processed first and only
// updated if > (not >=)
if (val_for_choice > max_val_for_choice)
{
max_val_for_choice = val_for_choice;
max_val = val;
expert = col + ii;
}
}
}
// Now, we perform the argmax reduce. We use the butterfly pattern so threads reach consensus about the max.
// This will be useful for K > 1 so that the threads can agree on "who" had the max value. That thread can
// then blank out their max with -inf and the warp can run more iterations...
#pragma unroll
for (int mask = THREADS_PER_ROW / 2; mask > 0; mask /= 2)
{
float other_max_for_choice = VLLM_SHFL_XOR_SYNC_WIDTH(max_val_for_choice, mask, THREADS_PER_ROW);
float other_max = VLLM_SHFL_XOR_SYNC_WIDTH(max_val, mask, THREADS_PER_ROW);
int other_expert = VLLM_SHFL_XOR_SYNC_WIDTH(expert, mask, THREADS_PER_ROW);
// We want lower indices to "win" in every thread so we break ties this way
if (other_max_for_choice > max_val_for_choice || (other_max_for_choice == max_val_for_choice && other_expert < expert))
{
max_val_for_choice = other_max_for_choice;
max_val = other_max;
expert = other_expert;
}
}
// Write the max for this k iteration to global memory.
if (thread_group_idx == 0)
{
// Add a guard to ignore experts not included by this node
const bool node_uses_expert = expert >= start_expert && expert < end_expert;
const bool should_process_row = row_is_active && node_uses_expert;
// The lead thread from each sub-group will write out the final results to global memory. (This will be a
// single) thread per row of the input/output matrices.
const int idx = k * thread_row + k_idx;
output[idx] = max_val;
indices[idx] = should_process_row ? (expert - start_expert) : NUM_EXPERTS;
source_rows[idx] = k_idx * num_rows + thread_row;
if (renormalize) {
selected_sum += max_val;
}
}
// Finally, we clear the value in the thread with the current max if there is another iteration to run.
if (k_idx + 1 < k)
{
const int ldg_group_for_expert = expert / COLS_PER_GROUP_LDG;
const int thread_to_clear_in_group = (expert / ELTS_PER_LDG) % THREADS_PER_ROW;
// Only the thread in the group which produced the max will reset the "winning" value to -inf.
if (thread_group_idx == thread_to_clear_in_group)
{
const int offset_for_expert = expert % ELTS_PER_LDG;
// Safe to set to any negative value since row_chunk values must be between 0 and 1.
row_chunk_for_choice[ldg_group_for_expert * ELTS_PER_LDG + offset_for_expert] = -10000.f;
}
}
}
// Renormalize the k weights for this row to sum to 1, if requested.
if (renormalize) {
if (thread_group_idx == 0)
{
const float denom = selected_sum > 0.f ? selected_sum : 1.f;
for (int k_idx = 0; k_idx < k; ++k_idx)
{
const int idx = k * thread_row + k_idx;
output[idx] = output[idx] / denom;
}
}
}
}
namespace detail
{
// Constructs some constants needed to partition the work across threads at compile time.
template <int EXPERTS, int BYTES_PER_LDG, int WARP_SIZE_PARAM, typename InputType>
struct TopkConstants
{
static constexpr int ELTS_PER_LDG = BYTES_PER_LDG / sizeof(InputType);
static_assert(EXPERTS / (ELTS_PER_LDG * WARP_SIZE_PARAM) == 0 || EXPERTS % (ELTS_PER_LDG * WARP_SIZE_PARAM) == 0, "");
static constexpr int VECs_PER_THREAD = MAX(1, EXPERTS / (ELTS_PER_LDG * WARP_SIZE_PARAM));
static constexpr int VPT = VECs_PER_THREAD * ELTS_PER_LDG;
static constexpr int THREADS_PER_ROW = EXPERTS / VPT;
static const int ROWS_PER_WARP = WARP_SIZE_PARAM / THREADS_PER_ROW;
};
} // namespace detail
template <int EXPERTS, int WARPS_PER_TB, int WARP_SIZE_PARAM, int MAX_BYTES_PER_LDG, typename IndType, typename InputType, ScoringFunc SF>
void topkGatingLauncherHelper(const InputType* input, const bool* finished, float* output, IndType* indices,
int* source_row, const int num_rows, const int k, const int start_expert, const int end_expert, const bool renormalize,
const float* bias, cudaStream_t stream)
{
static constexpr int BYTES_PER_LDG = MIN(MAX_BYTES_PER_LDG, sizeof(InputType) * EXPERTS);
using Constants = detail::TopkConstants<EXPERTS, BYTES_PER_LDG, WARP_SIZE_PARAM, InputType>;
static constexpr int VPT = Constants::VPT;
static constexpr int ROWS_PER_WARP = Constants::ROWS_PER_WARP;
const int num_warps = (num_rows + ROWS_PER_WARP - 1) / ROWS_PER_WARP;
const int num_blocks = (num_warps + WARPS_PER_TB - 1) / WARPS_PER_TB;
dim3 block_dim(WARP_SIZE_PARAM, WARPS_PER_TB);
topkGating<VPT, EXPERTS, WARPS_PER_TB, BYTES_PER_LDG, WARP_SIZE_PARAM, IndType, InputType, SF><<<num_blocks, block_dim, 0, stream>>>(
input, finished, output, num_rows, indices, source_row, k, start_expert, end_expert, renormalize, bias);
}
#ifndef USE_ROCM
#define LAUNCH_TOPK(NUM_EXPERTS, WARPS_PER_TB, MAX_BYTES) \
static_assert(WARP_SIZE == 32, \
"Unsupported warp size. Only 32 is supported for CUDA"); \
topkGatingLauncherHelper<NUM_EXPERTS, WARPS_PER_TB, WARP_SIZE, MAX_BYTES, \
IndType, InputType, SF>( \
gating_output, nullptr, topk_weights, topk_indices, \
token_expert_indices, num_tokens, topk, 0, num_experts, renormalize, \
bias, stream);
#else
#define LAUNCH_TOPK(NUM_EXPERTS, WARPS_PER_TB, MAX_BYTES) \
if (WARP_SIZE == 64) { \
topkGatingLauncherHelper<NUM_EXPERTS, WARPS_PER_TB, 64, MAX_BYTES, \
IndType, InputType, SF>( \
gating_output, nullptr, topk_weights, topk_indices, \
token_expert_indices, num_tokens, topk, 0, num_experts, renormalize, \
bias, stream); \
} else if (WARP_SIZE == 32) { \
topkGatingLauncherHelper<NUM_EXPERTS, WARPS_PER_TB, 32, MAX_BYTES, \
IndType, InputType, SF>( \
gating_output, nullptr, topk_weights, topk_indices, \
token_expert_indices, num_tokens, topk, 0, num_experts, renormalize, \
bias, stream); \
} else { \
assert(false && \
"Unsupported warp size. Only 32 and 64 are supported for ROCm"); \
}
#endif
template <typename IndType, typename InputType, ScoringFunc SF>
void topkGatingKernelLauncher(
const InputType* gating_output,
float* topk_weights,
IndType* topk_indices,
int* token_expert_indices,
float* workspace,
const int num_tokens,
const int num_experts,
const int topk,
const bool renormalize,
const float* bias,
cudaStream_t stream) {
static constexpr int WARPS_PER_TB = 4;
static constexpr int BYTES_PER_LDG_POWER_OF_2 = 16;
#ifndef USE_ROCM
// for bfloat16 dtype, we need 4 bytes loading to make sure num_experts
// elements can be loaded by a warp
static constexpr int BYTES_PER_LDG_MULTIPLE_64 =
(std::is_same_v<InputType, __nv_bfloat16> || std::is_same_v<InputType, __half>) ? 4 : 8;
#endif
switch (num_experts) {
case 1:
LAUNCH_TOPK(1, WARPS_PER_TB, BYTES_PER_LDG_POWER_OF_2);
break;
case 2:
LAUNCH_TOPK(2, WARPS_PER_TB, BYTES_PER_LDG_POWER_OF_2);
break;
case 4:
LAUNCH_TOPK(4, WARPS_PER_TB, BYTES_PER_LDG_POWER_OF_2);
break;
case 8:
LAUNCH_TOPK(8, WARPS_PER_TB, BYTES_PER_LDG_POWER_OF_2);
break;
case 16:
LAUNCH_TOPK(16, WARPS_PER_TB, BYTES_PER_LDG_POWER_OF_2);
break;
case 32:
LAUNCH_TOPK(32, WARPS_PER_TB, BYTES_PER_LDG_POWER_OF_2);
break;
case 64:
LAUNCH_TOPK(64, WARPS_PER_TB, BYTES_PER_LDG_POWER_OF_2);
break;
case 128:
LAUNCH_TOPK(128, WARPS_PER_TB, BYTES_PER_LDG_POWER_OF_2);
break;
case 256:
LAUNCH_TOPK(256, WARPS_PER_TB, BYTES_PER_LDG_POWER_OF_2);
break;
case 512:
LAUNCH_TOPK(512, WARPS_PER_TB, BYTES_PER_LDG_POWER_OF_2);
break;
// (CUDA only) support multiples of 64 when num_experts is not power of 2.
// ROCm uses WARP_SIZE 64 so 8 bytes loading won't fit for some of num_experts,
// alternatively we can test 4 bytes loading and enable it in future.
#ifndef USE_ROCM
case 192:
LAUNCH_TOPK(192, WARPS_PER_TB, BYTES_PER_LDG_MULTIPLE_64);
break;
case 320:
LAUNCH_TOPK(320, WARPS_PER_TB, BYTES_PER_LDG_MULTIPLE_64);
break;
case 384:
LAUNCH_TOPK(384, WARPS_PER_TB, BYTES_PER_LDG_MULTIPLE_64);
break;
case 448:
LAUNCH_TOPK(448, WARPS_PER_TB, BYTES_PER_LDG_MULTIPLE_64);
break;
case 576:
LAUNCH_TOPK(576, WARPS_PER_TB, BYTES_PER_LDG_MULTIPLE_64);
break;
#endif
default: {
STD_TORCH_CHECK(workspace != nullptr,
"workspace must be provided for num_experts that are not a power of 2 or multiple of 64.");
static constexpr int TPB = 256;
if constexpr (SF == SCORING_SOFTMAX) {
moeSoftmax<TPB, InputType><<<num_tokens, TPB, 0, stream>>>(
gating_output, nullptr, workspace, num_experts);
} else if constexpr (SF == SCORING_SIGMOID) {
moeSigmoid<TPB, InputType><<<num_tokens, TPB, 0, stream>>>(
gating_output, nullptr, workspace, num_experts);
} else {
STD_TORCH_CHECK(false, "Unsupported scoring func");
}
moeTopK<TPB><<<num_tokens, TPB, 0, stream>>>(
workspace, nullptr, topk_weights, topk_indices, token_expert_indices,
num_experts, topk, 0, num_experts, renormalize, bias);
}
}
}
} // namespace moe
} // namespace vllm
template<typename ComputeType, vllm::moe::ScoringFunc SF>
void dispatch_topk_launch(
torch::stable::Tensor& gating_output,
torch::stable::Tensor& topk_weights,
torch::stable::Tensor& topk_indices,
torch::stable::Tensor& token_expert_indices,
torch::stable::Tensor& softmax_workspace,
int num_tokens, int num_experts, int topk, bool renormalize,
std::optional<torch::stable::Tensor> bias,
cudaStream_t stream)
{
const float* bias_ptr = nullptr;
if (bias.has_value()) {
const torch::stable::Tensor& bias_tensor = bias.value();
STD_TORCH_CHECK(bias_tensor.scalar_type() == torch::headeronly::ScalarType::Float,
"bias tensor must be float32");
STD_TORCH_CHECK(bias_tensor.dim() == 1, "bias tensor must be 1D");
STD_TORCH_CHECK(bias_tensor.size(0) == num_experts,
"bias size mismatch, expected: ", num_experts);
STD_TORCH_CHECK(bias_tensor.is_contiguous(), "bias tensor must be contiguous");
bias_ptr = bias_tensor.const_data_ptr<float>();
}
if (topk_indices.scalar_type() == torch::headeronly::ScalarType::Int) {
vllm::moe::topkGatingKernelLauncher<int, ComputeType, SF>(
reinterpret_cast<const ComputeType*>(gating_output.const_data_ptr()),
topk_weights.mutable_data_ptr<float>(),
topk_indices.mutable_data_ptr<int>(),
token_expert_indices.mutable_data_ptr<int>(),
softmax_workspace.mutable_data_ptr<float>(),
num_tokens, num_experts, topk, renormalize,
bias_ptr, stream);
} else if (topk_indices.scalar_type() == torch::headeronly::ScalarType::UInt32) {
vllm::moe::topkGatingKernelLauncher<uint32_t, ComputeType, SF>(
reinterpret_cast<const ComputeType*>(gating_output.const_data_ptr()),
topk_weights.mutable_data_ptr<float>(),
topk_indices.mutable_data_ptr<uint32_t>(),
token_expert_indices.mutable_data_ptr<int>(),
softmax_workspace.mutable_data_ptr<float>(),
num_tokens, num_experts, topk, renormalize,
bias_ptr, stream);
} else {
STD_TORCH_CHECK(topk_indices.scalar_type() == torch::headeronly::ScalarType::Long);
vllm::moe::topkGatingKernelLauncher<int64_t, ComputeType, SF>(
reinterpret_cast<const ComputeType*>(gating_output.const_data_ptr()),
topk_weights.mutable_data_ptr<float>(),
topk_indices.mutable_data_ptr<int64_t>(),
token_expert_indices.mutable_data_ptr<int>(),
softmax_workspace.mutable_data_ptr<float>(),
num_tokens, num_experts, topk, renormalize,
bias_ptr, stream);
}
}
void topk_softmax(
torch::stable::Tensor& topk_weights, // [num_tokens, topk]
torch::stable::Tensor& topk_indices, // [num_tokens, topk]
torch::stable::Tensor& token_expert_indices, // [num_tokens, topk]
torch::stable::Tensor& gating_output, // [num_tokens, num_experts]
bool renormalize,
std::optional<torch::stable::Tensor> bias)
{
const int num_experts = gating_output.size(-1);
const auto num_tokens = gating_output.numel() / num_experts;
const int topk = topk_weights.size(-1);
const bool is_pow_2 = (num_experts != 0) && ((num_experts & (num_experts - 1)) == 0);
const bool needs_workspace = !is_pow_2 || num_experts > 256;
const int64_t workspace_size = needs_workspace ? num_tokens * num_experts : 0;
torch::stable::accelerator::DeviceGuard guard(gating_output.get_device_index());
const cudaStream_t stream =
get_current_cuda_stream(gating_output.get_device_index());
auto softmax_workspace = torch::stable::new_empty(
gating_output, {workspace_size}, torch::headeronly::ScalarType::Float);
if (gating_output.scalar_type() == torch::headeronly::ScalarType::Float) {
dispatch_topk_launch<float, vllm::moe::SCORING_SOFTMAX>(gating_output, topk_weights, topk_indices,
token_expert_indices, softmax_workspace, num_tokens, num_experts, topk, renormalize,
bias, stream);
} else if (gating_output.scalar_type() == torch::headeronly::ScalarType::Half) {
dispatch_topk_launch<__half, vllm::moe::SCORING_SOFTMAX>(gating_output, topk_weights, topk_indices,
token_expert_indices, softmax_workspace, num_tokens, num_experts, topk, renormalize,
bias, stream);
} else if (gating_output.scalar_type() == torch::headeronly::ScalarType::BFloat16) {
dispatch_topk_launch<__nv_bfloat16, vllm::moe::SCORING_SOFTMAX>(gating_output, topk_weights, topk_indices,
token_expert_indices, softmax_workspace, num_tokens, num_experts, topk, renormalize,
bias, stream);
} else {
STD_TORCH_CHECK(false, "Unsupported gating_output data type: ", gating_output.scalar_type());
}
}
void topk_sigmoid(
torch::stable::Tensor& topk_weights, // [num_tokens, topk]
torch::stable::Tensor& topk_indices, // [num_tokens, topk]
torch::stable::Tensor& token_expert_indices, // [num_tokens, topk]
torch::stable::Tensor& gating_output, // [num_tokens, num_experts]
bool renormalize,
std::optional<torch::stable::Tensor> bias)
{
const int num_experts = gating_output.size(-1);
const auto num_tokens = gating_output.numel() / num_experts;
const int topk = topk_weights.size(-1);
const bool is_pow_2 = (num_experts != 0) && ((num_experts & (num_experts - 1)) == 0);
const bool needs_workspace = !is_pow_2 || num_experts > 256;
const int64_t workspace_size = needs_workspace ? num_tokens * num_experts : 0;
torch::stable::accelerator::DeviceGuard guard(gating_output.get_device_index());
const cudaStream_t stream =
get_current_cuda_stream(gating_output.get_device_index());
auto workspace = torch::stable::new_empty(
gating_output, {workspace_size}, torch::headeronly::ScalarType::Float);
if (gating_output.scalar_type() == torch::headeronly::ScalarType::Float) {
dispatch_topk_launch<float, vllm::moe::SCORING_SIGMOID>(gating_output, topk_weights, topk_indices,
token_expert_indices, workspace, num_tokens, num_experts, topk, renormalize,
bias, stream);
} else if (gating_output.scalar_type() == torch::headeronly::ScalarType::Half) {
dispatch_topk_launch<__half, vllm::moe::SCORING_SIGMOID>(gating_output, topk_weights, topk_indices,
token_expert_indices, workspace, num_tokens, num_experts, topk, renormalize,
bias, stream);
} else if (gating_output.scalar_type() == torch::headeronly::ScalarType::BFloat16) {
dispatch_topk_launch<__nv_bfloat16, vllm::moe::SCORING_SIGMOID>(gating_output, topk_weights, topk_indices,
token_expert_indices, workspace, num_tokens, num_experts, topk, renormalize,
bias, stream);
} else {
STD_TORCH_CHECK(false, "Unsupported gating_output data type: ", gating_output.scalar_type());
}
}

124
upstream_ref/xllm_latest/core/kernels/cuda/moe/fused_moe.cpp Normal file
View File

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

105
upstream_ref/xllm_latest/core/kernels/cuda/moe/moe_combine.cu Executable file
View File

@@ -0,0 +1,105 @@
/* Copyright 2025-2026 The xLLM Authors. All Rights Reserved.
Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at
https://github.com/jd-opensource/xllm/blob/main/LICENSE
Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
==============================================================================*/
// Fused MoE combine kernel — reorder + weighted sum in one pass.
// Replaces: torch::zeros + index_copy_ + view + multiply + sum
//
// Algorithm per token (each block handles one token):
// 1. For each of its topk experts, read gemm2 at flat_idx directly
// (gemm2 is flat-index-ordered after scatter via index_copy_ with dst_src)
// 2. Multiply by router weight
// 3. Accumulate into output[token]
//
// Grid: num_tokens (N) blocks
// Block: HIDDEN_DIM / HIDDEN_TILE threads
#include <c10/cuda/CUDAGuard.h>
#include "device_utils.cuh"
#include "kernels/cuda/cuda_ops_api.h"
namespace xllm::kernel::cuda {
constexpr int32_t kCombineBlockSize = 256;
template <typename scalar_t>
__global__ void XLLM_KERNEL_ATTR(kCombineBlockSize) moe_combine_kernel(
const scalar_t* __restrict__ gemm2, // [N*topk, H] flat-index-ordered
const float* __restrict__ reduce_weight, // [N, topk]
scalar_t* __restrict__ output, // [N, H]
int64_t N,
int32_t topk,
int64_t H) {
int64_t token_id = blockIdx.x; // 0 .. N-1
if (token_id >= N) return;
int32_t tid = threadIdx.x;
int32_t stride = kCombineBlockSize;
// Accumulate over topk experts for this token
for (int64_t h = tid; h < H; h += stride) {
float acc = 0.0f;
for (int32_t k = 0; k < topk; ++k) {
int64_t flat_idx = token_id * topk + k;
float w = reduce_weight[flat_idx];
acc += w * static_cast<float>(gemm2[flat_idx * H + h]);
}
output[token_id * H + h] = static_cast<scalar_t>(acc);
}
}
// ---- Host-side orchestrator ----
torch::Tensor moe_combine_result(
const torch::Tensor& gemm2, // [N*topk, H] flat-index-ordered
const torch::Tensor& reduce_weight, // [N, topk] float or same as gemm2
int64_t N,
int32_t topk) {
auto stream = at::cuda::getCurrentCUDAStream();
int64_t H = gemm2.size(1);
auto dtype = gemm2.scalar_type();
auto output = torch::empty({N, H}, gemm2.options());
auto rw = reduce_weight.to(gemm2.device(), torch::kFloat32).contiguous();
if (dtype == torch::kFloat16) {
moe_combine_kernel<c10::Half>
<<<N, kCombineBlockSize, 0, stream>>>(gemm2.data_ptr<c10::Half>(),
rw.data_ptr<float>(),
output.data_ptr<c10::Half>(),
N,
topk,
H);
} else if (dtype == torch::kBFloat16) {
moe_combine_kernel<c10::BFloat16>
<<<N, kCombineBlockSize, 0, stream>>>(gemm2.data_ptr<c10::BFloat16>(),
rw.data_ptr<float>(),
output.data_ptr<c10::BFloat16>(),
N,
topk,
H);
} else {
moe_combine_kernel<float>
<<<N, kCombineBlockSize, 0, stream>>>(gemm2.data_ptr<float>(),
rw.data_ptr<float>(),
output.data_ptr<float>(),
N,
topk,
H);
}
return output;
}
} // namespace xllm::kernel::cuda

155
upstream_ref/xllm_latest/core/kernels/cuda/moe/moe_compute_index.cu Normal file
View File

@@ -0,0 +1,155 @@
/* Copyright 2025-2026 The xLLM Authors. All Rights Reserved.
Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at
https://github.com/jd-opensource/xllm/blob/main/LICENSE
Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
==============================================================================*/
// Fused MoE token index computation — 3 kernels replacing:
// torch::bincount + 2 × torch::argsort + torch::cumsum + CPU sync
//
// Phase 1 histogram: atomicAdd per-expert token counts
// Phase 2 prefix_sum: 1 block, exclusive scan → expert_offsets
// Phase 3 place_indices: atomicAdd on offsets, write dst_src + src_dst
//
// expert_sizes = per-expert token count [num_experts] (preserved)
// expert_offsets = exclusive prefix sum of counts (scratch, reused)
#include <c10/cuda/CUDAGuard.h>
#include <cub/block/block_scan.cuh>
#include "kernels/cuda/cuda_ops_api.h"
namespace xllm::kernel::cuda {
constexpr int32_t kMoeIndexBlock = 256;
// ---- Phase 1: histogram ----
__global__ void
#ifdef USE_DCU
__launch_bounds__(kMoeIndexBlock, 1)
#endif
moe_histogram_kernel(const int32_t* __restrict__ expert_id,
int32_t* __restrict__ expert_sizes,
int64_t num_elements,
int32_t num_experts) {
int64_t tid = int64_t(blockIdx.x) * kMoeIndexBlock + threadIdx.x;
if (tid < num_elements) {
int32_t eid = expert_id[tid];
if (eid >= 0 && eid < num_experts) {
atomicAdd(&expert_sizes[eid], 1);
}
}
}
// ---- Phase 2: exclusive prefix sum (1 block) ----
// input: expert_sizes (per-expert counts)
// output: expert_offsets (exclusive scan of counts)
// total_out (total number of tokens, scalar)
__global__ void
#ifdef USE_DCU
__launch_bounds__(kMoeIndexBlock, 1)
#endif
moe_prefix_sum_kernel(const int32_t* __restrict__ expert_sizes,
int32_t* __restrict__ expert_offsets,
int32_t num_experts,
int64_t* __restrict__ total_out) {
using BlockScan = cub::BlockScan<int32_t, kMoeIndexBlock>;
__shared__ typename BlockScan::TempStorage s_scan;
int32_t val = (threadIdx.x < num_experts) ? expert_sizes[threadIdx.x] : 0;
int32_t offset;
BlockScan(s_scan).ExclusiveSum(val, offset);
__syncthreads();
// total = all elements sum = last thread's exclusive output + its input
int32_t total = offset + val;
if (threadIdx.x < num_experts) {
expert_offsets[threadIdx.x] = offset;
}
if (threadIdx.x == 0 && total_out != nullptr) {
*total_out = total;
}
}
// ---- Phase 3: place indices ----
// atomicAdd on expert_offsets to assign a unique position within
// [start(e), start(e)+count(e)), then write both direction mappings.
__global__ void
#ifdef USE_DCU
__launch_bounds__(kMoeIndexBlock, 1)
#endif
moe_place_indices_kernel(const int32_t* __restrict__ expert_id,
int32_t* __restrict__ expert_offsets,
int32_t* __restrict__ dst_src,
int32_t* __restrict__ src_dst,
int64_t num_elements,
int32_t num_experts) {
int64_t flat_idx = int64_t(blockIdx.x) * kMoeIndexBlock + threadIdx.x;
if (flat_idx >= num_elements) return;
int32_t eid = expert_id[flat_idx];
if (eid < 0 || eid >= num_experts) return;
int32_t pos = atomicAdd(&expert_offsets[eid], 1);
dst_src[pos] = static_cast<int32_t>(flat_idx);
src_dst[flat_idx] = pos;
}
// ---- Host-side orchestrator ----
// Returns {src_dst, dst_src, expert_sizes}
std::tuple<torch::Tensor, torch::Tensor, torch::Tensor> moe_compute_index(
const torch::Tensor& expert_id,
int64_t num_experts) {
auto device = expert_id.device();
auto stream = at::cuda::getCurrentCUDAStream();
int64_t N = expert_id.numel();
int32_t E = static_cast<int32_t>(num_experts);
CHECK_LE(E, kMoeIndexBlock) << "num_experts cannot exceed " << kMoeIndexBlock;
auto expert_id_i32 = expert_id.to(torch::kInt32).contiguous();
auto opt_i32 = expert_id_i32.options();
auto expert_sizes = torch::zeros({num_experts}, opt_i32);
auto expert_offsets = torch::empty({num_experts}, opt_i32);
auto dst_src = torch::empty({N}, opt_i32);
auto src_dst = torch::empty({N}, opt_i32);
int64_t grid = (N + kMoeIndexBlock - 1) / kMoeIndexBlock;
// Phase 1: histogram
moe_histogram_kernel<<<grid, kMoeIndexBlock, 0, stream>>>(
expert_id_i32.data_ptr<int32_t>(),
expert_sizes.data_ptr<int32_t>(),
N,
E);
// Phase 2: prefix sum (1 block)
moe_prefix_sum_kernel<<<1, kMoeIndexBlock, 0, stream>>>(
expert_sizes.data_ptr<int32_t>(),
expert_offsets.data_ptr<int32_t>(),
E,
nullptr);
// Phase 3: place indices
moe_place_indices_kernel<<<grid, kMoeIndexBlock, 0, stream>>>(
expert_id_i32.data_ptr<int32_t>(),
expert_offsets.data_ptr<int32_t>(),
dst_src.data_ptr<int32_t>(),
src_dst.data_ptr<int32_t>(),
N,
E);
return std::make_tuple(src_dst, dst_src, expert_sizes);
}
} // namespace xllm::kernel::cuda

59
upstream_ref/xllm_latest/core/kernels/cuda/moe/moe_fused_topk.cu Normal file
View File

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

345
upstream_ref/xllm_latest/core/kernels/cuda/moe/moe_topk.cuh Normal file
View File

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

609
upstream_ref/xllm_latest/core/kernels/cuda/moe/moe_topk_sigmoid_kernels.cuh Normal file
View File

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

867
upstream_ref/xllm_latest/core/kernels/cuda/moe/moe_topk_softmax_kernels.cuh Normal file
View File

@@ -0,0 +1,867 @@
// Adapt from
// https://github.com/vllm-project/vllm/blob/v0.7.3/csrc/moe/topk_softmax_kernels.cu
// which is originally adapted from
// https://github.com/NVIDIA/TensorRT-LLM/blob/v0.7.1/cpp/tensorrt_llm/kernels/mixtureOfExperts/moe_kernels.cu
/* Copyright 2025 SGLang Team. All Rights Reserved.
Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at
http://www.apache.org/licenses/LICENSE-2.0
Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
==============================================================================*/
#include <ATen/cuda/CUDAContext.h>
#include <c10/cuda/CUDAGuard.h>
#include <torch/all.h>
#include <cub/util_type.cuh>
#if !defined(USE_DCU) && !defined(USE_MACA)
#include <cuda/functional>
#endif
#include "kernels/cuda/device_utils.cuh"
using cub_kvp = cub::KeyValuePair<int, float>;
namespace {
using namespace xllm::kernel::cuda;
#if defined(USE_DCU)
static constexpr unsigned long long kSoftmaxFullMask = 0xffffffffffffffffULL;
#else
static constexpr unsigned int kSoftmaxFullMask = 0xffffffffU;
#endif
// ====================== Softmax things ===============================
// We have our own implementation of softmax here so we can support transposing
// the output in the softmax kernel when we extend this module to support
// expert-choice routing.
template <typename T, int TPB>
__launch_bounds__(TPB) __global__
void moe_softmax(const T* input,
const bool* finished,
float* output,
const int num_cols,
const float moe_softcapping,
const float* correction_bias) {
using BlockReduce = cub::BlockReduce<float, TPB>;
__shared__ typename BlockReduce::TempStorage tmpStorage;
__shared__ float normalizing_factor;
__shared__ float float_max;
const int thread_row_offset = blockIdx.x * num_cols;
float threadData(-FLT_MAX);
// Don't touch finished rows.
if ((finished != nullptr) && finished[blockIdx.x]) {
return;
}
// First pass: Apply transformation, find max, and write transformed values to
// output
for (int ii = threadIdx.x; ii < num_cols; ii += TPB) {
const int idx = thread_row_offset + ii;
float val = convert_to_float<T>(input[idx]);
// Apply tanh softcapping if enabled
if (moe_softcapping != 0.0f) {
val = tanhf(val / moe_softcapping) * moe_softcapping;
}
// Apply correction bias if provided
if (correction_bias != nullptr) {
val = val + correction_bias[ii];
}
output[idx] = val; // Store transformed value
threadData = max(val, threadData);
}
const float maxElem =
BlockReduce(tmpStorage).Reduce(threadData, MaxReduceOp());
if (threadIdx.x == 0) {
float_max = maxElem;
}
__syncthreads();
// Second pass: Compute sum using transformed values from output
threadData = 0;
for (int ii = threadIdx.x; ii < num_cols; ii += TPB) {
const int idx = thread_row_offset + ii;
threadData += exp((output[idx] - float_max));
}
const auto Z = BlockReduce(tmpStorage).Sum(threadData);
if (threadIdx.x == 0) {
normalizing_factor = 1.f / Z;
}
__syncthreads();
// Third pass: Compute final softmax using transformed values from output
for (int ii = threadIdx.x; ii < num_cols; ii += TPB) {
const int idx = thread_row_offset + ii;
const float softmax_val =
exp((output[idx] - float_max)) * normalizing_factor;
output[idx] = softmax_val;
}
}
namespace moe {
class TopKPair {
public:
static constexpr int kPair = 2;
static constexpr int kMaxIndex = 0;
cub_kvp max;
cub_kvp secondMax;
__device__ TopKPair() {}
__device__ TopKPair(cub_kvp max, cub_kvp secondMax)
: max(max), secondMax(secondMax) {}
};
class TopKPairArgMax {
public:
__device__ TopKPairArgMax() {}
__device__ __forceinline__ TopKPair
operator()(const TopKPair& candidate1, const TopKPair& candidate2) const {
cub_kvp globalMax, globalSecondMax;
// Determine the global maximum
if (candidate1.max.value > candidate2.max.value) {
globalMax = candidate1.max;
} else {
globalMax = candidate2.max;
}
// Determine the global second maximum
if (globalMax.key == candidate1.max.key) {
// If candidate1 contributed the max, compare its secondMax with
// candidate2's max
globalSecondMax = (candidate1.secondMax.value > candidate2.max.value)
? candidate1.secondMax
: candidate2.max;
} else {
// If candidate2 contributed the max, compare its secondMax with
// candidate1's max
globalSecondMax = (candidate2.secondMax.value > candidate1.max.value)
? candidate2.secondMax
: candidate1.max;
}
return TopKPair(globalMax, globalSecondMax);
}
};
} // namespace moe
template <int TPB>
__launch_bounds__(TPB) __global__
void moe_topk_fast(float* inputs_after_softmax,
const bool* finished,
float* output,
int* indices,
const int num_experts,
const int k,
const int start_expert,
const int end_expert,
const bool renormalize) {
using namespace moe;
using BlockReduce = cub::BlockReduce<TopKPair, TPB>;
__shared__ typename BlockReduce::TempStorage tmpStorage;
TopKPair thread_pair;
const int block_row = blockIdx.x;
const bool row_is_active = finished ? !finished[block_row] : true;
const int thread_read_offset = blockIdx.x * num_experts;
float row_sum_for_renormalize = 0;
// Each loop finds the top 2 elements,
// thus requiring only ceil(k / 2) loops (calculated as (k + 1) / 2).
for (int k_idx = 0; k_idx < (k + TopKPair::kPair - 1) / TopKPair::kPair;
++k_idx) {
// Initializing the top 2 elements by the minimum value.
thread_pair.max.key = 0;
thread_pair.max.value = -1.f;
thread_pair.secondMax.key = 0;
thread_pair.secondMax.value = -1.f;
cub_kvp inp_kvp;
for (int expert = threadIdx.x; expert < num_experts; expert += TPB) {
const int idx = thread_read_offset + expert;
inp_kvp.key = expert;
inp_kvp.value = inputs_after_softmax[idx];
// updating the thread_pair according to inp_kvp's value
if (inp_kvp.value > thread_pair.max.value) {
thread_pair.secondMax = thread_pair.max;
thread_pair.max = inp_kvp;
} else if (inp_kvp.value > thread_pair.secondMax.value) {
thread_pair.secondMax = inp_kvp;
}
}
TopKPairArgMax reducer;
const TopKPair result_pair =
BlockReduce(tmpStorage).Reduce(thread_pair, reducer);
if (threadIdx.x == 0) {
#pragma unroll
// updating 2 elements to the result.
for (int i = 0; i < TopKPair::kPair; i++) {
if (k_idx * 2 + i >= k) {
break;
}
cub_kvp result = (i == TopKPair::kMaxIndex) ? result_pair.max
: result_pair.secondMax;
int expert = result.key;
bool node_uses_expert = expert >= start_expert && expert < end_expert;
bool should_process_row = row_is_active && node_uses_expert;
// The inputs_after_softmax is modified in-place to avoid unnecessary
// loops for finding the top k-1 value. 1.f represents the minimum
// value.
inputs_after_softmax[thread_read_offset + expert] = -1.f;
int idx = k * block_row + k_idx * 2 + i;
output[idx] = result.value;
indices[idx] =
should_process_row ? (expert - start_expert) : num_experts;
assert(indices[idx] >= 0);
row_sum_for_renormalize += result.value;
}
}
__syncthreads();
}
if (renormalize && threadIdx.x == 0) {
float row_sum_for_renormalize_inv = 1.f / row_sum_for_renormalize;
for (int k_idx = 0; k_idx < k; ++k_idx) {
const int idx = k * block_row + k_idx;
output[idx] = output[idx] * row_sum_for_renormalize_inv;
}
}
}
template <int TPB>
__launch_bounds__(TPB) __global__ void moe_topK(float* inputs_after_softmax,
const bool* finished,
float* output,
int* indices,
const int num_experts,
const int k,
const int start_expert,
const int end_expert,
const bool renormalize) {
using cub_kvp = cub::KeyValuePair<int, float>;
using BlockReduce = cub::BlockReduce<cub_kvp, TPB>;
__shared__ typename BlockReduce::TempStorage tmpStorage;
cub_kvp thread_kvp;
cub::ArgMax arg_max;
const int block_row = blockIdx.x;
const bool row_is_active = finished ? !finished[block_row] : true;
const int thread_read_offset = blockIdx.x * num_experts;
float row_sum_for_renormalize = 0;
for (int k_idx = 0; k_idx < k; ++k_idx) {
thread_kvp.key = 0;
thread_kvp.value = -1.f; // This is OK because inputs are probabilities
cub_kvp inp_kvp;
for (int expert = threadIdx.x; expert < num_experts; expert += TPB) {
const int idx = thread_read_offset + expert;
inp_kvp.key = expert;
inp_kvp.value = inputs_after_softmax[idx];
thread_kvp = arg_max(inp_kvp, thread_kvp);
}
const cub_kvp result_kvp =
BlockReduce(tmpStorage).Reduce(thread_kvp, arg_max);
if (threadIdx.x == 0) {
// Ignore experts the node isn't responsible for with expert parallelism
const int expert = result_kvp.key;
const bool node_uses_expert =
expert >= start_expert && expert < end_expert;
const bool should_process_row = row_is_active && node_uses_expert;
const int idx = k * block_row + k_idx;
output[idx] = result_kvp.value;
indices[idx] = should_process_row ? (expert - start_expert) : num_experts;
assert(indices[idx] >= 0);
row_sum_for_renormalize += result_kvp.value;
// The inputs_after_softmax is modified in-place to avoid unnecessary
// loops for finding the top k-1 value. 1.f represents the minimum value.
inputs_after_softmax[thread_read_offset + expert] = -1.f;
}
__syncthreads();
}
if (renormalize && threadIdx.x == 0) {
float row_sum_for_renormalize_inv = 1.f / row_sum_for_renormalize;
for (int k_idx = 0; k_idx < k; ++k_idx) {
const int idx = k * block_row + k_idx;
output[idx] = output[idx] * row_sum_for_renormalize_inv;
}
}
}
// ====================== TopK softmax things ===============================
/*
A Top-K gating softmax written to exploit when the number of experts in the
MoE layers are a small power of 2. This allows us to cleanly share the rows
among the threads in a single warp and eliminate communication between warps
(so no need to use shared mem).
It fuses the softmax, max and argmax into a single kernel.
Limitations:
1) This implementation is intended for when the number of experts is a small
power of 2. 2) This implementation assumes k is small, but will work for any
k.
*/
template <typename T,
int VPT,
int NUM_EXPERTS,
int WARPS_PER_CTA,
int BYTES_PER_LDG>
__launch_bounds__(WARPS_PER_CTA* WARP_SIZE) __global__
void topk_gating_softmax(const T* input,
const bool* finished,
float* output,
const int num_rows,
int* indices,
const int k,
const int start_expert,
const int end_expert,
const bool renormalize,
const float moe_softcapping,
const float* correction_bias) {
// We begin by enforcing compile time assertions and setting up compile time
// constants.
static_assert(VPT == (VPT & -VPT), "VPT must be power of 2");
static_assert(NUM_EXPERTS == (NUM_EXPERTS & -NUM_EXPERTS),
"NUM_EXPERTS must be power of 2");
static_assert(BYTES_PER_LDG == (BYTES_PER_LDG & -BYTES_PER_LDG),
"BYTES_PER_LDG must be power of 2");
static_assert(BYTES_PER_LDG <= 16, "BYTES_PER_LDG must be leq 16");
// Number of bytes each thread pulls in per load
static constexpr int kEltsPerLdg = BYTES_PER_LDG / sizeof(T);
static constexpr int kEltsPerRow = NUM_EXPERTS;
static constexpr int kThreadsPerRow = kEltsPerRow / VPT;
static constexpr int kLdgPerThread = VPT / kEltsPerLdg;
// Restrictions based on previous section.
static_assert(
VPT % kEltsPerLdg == 0,
"The elements per thread must be a multiple of the elements per ldg");
static_assert(WARP_SIZE % kThreadsPerRow == 0,
"The threads per row must cleanly divide the threads per warp");
static_assert(kThreadsPerRow == (kThreadsPerRow & -kThreadsPerRow),
"THREADS_PER_ROW must be power of 2");
static_assert(kThreadsPerRow <= WARP_SIZE,
"THREADS_PER_ROW can be at most warp size");
// We have NUM_EXPERTS elements per row. We specialize for small #experts
static constexpr int kEltsPerWarp = WARP_SIZE * VPT;
static constexpr int kRowsPerWarp = kEltsPerWarp / kEltsPerRow;
static constexpr int kRowsPerCta = WARPS_PER_CTA * kRowsPerWarp;
// Restrictions for previous section.
static_assert(kEltsPerWarp % kEltsPerRow == 0,
"The elts per row must cleanly divide the total elt per warp");
// ===================== From this point, we finally start computing run-time
// variables. ========================
// Compute CTA and warp rows. We pack multiple rows into a single warp, and a
// block contains WARPS_PER_CTA warps. This, each block processes a chunk of
// rows. We start by computing the start row for each block.
const int cta_base_row = blockIdx.x * kRowsPerCta;
// Now, using the base row per thread block, we compute the base row per warp.
const int warp_base_row = cta_base_row + threadIdx.y * kRowsPerWarp;
// The threads in a warp are split into sub-groups that will work on a row.
// We compute row offset for each thread sub-group
const int thread_row_in_warp = threadIdx.x / kThreadsPerRow;
const int thread_row = warp_base_row + thread_row_in_warp;
// Threads with indices out of bounds should early exit here.
if (thread_row >= num_rows) {
return;
}
const bool row_is_active = finished ? !finished[thread_row] : true;
// We finally start setting up the read pointers for each thread. First, each
// thread jumps to the start of the row it will read.
const T* thread_row_ptr = input + thread_row * kEltsPerRow;
// Now, we compute the group each thread belong to in order to determine the
// first column to start loads.
const int thread_group_idx = threadIdx.x % kThreadsPerRow;
const int first_elt_read_by_thread = thread_group_idx * kEltsPerLdg;
const T* thread_read_ptr = thread_row_ptr + first_elt_read_by_thread;
// Determine the pointer type to use to read in the data depending on the
// BYTES_PER_LDG template param. In theory, this can support all powers of 2
// up to 16. NOTE(woosuk): The original implementation uses CUTLASS aligned
// array here. We defined our own aligned array and use it here to avoid the
// dependency on CUTLASS.
using AccessType = AlignedArray<T, kEltsPerLdg>;
// Finally, we pull in the data from global mem
T row_chunk_temp[VPT];
AccessType* row_chunk_vec_ptr =
reinterpret_cast<AccessType*>(&row_chunk_temp);
const AccessType* vec_thread_read_ptr =
reinterpret_cast<const AccessType*>(thread_read_ptr);
#pragma unroll
// Note(Byron): interleaved loads to achieve better memory coalescing
// | thread[0] | thread[1] | thread[2] | thread[3] | thread[0] | thread[1] |
// thread[2] | thread[3] | ...
for (int ii = 0; ii < kLdgPerThread; ++ii) {
row_chunk_vec_ptr[ii] = vec_thread_read_ptr[ii * kThreadsPerRow];
}
float row_chunk[VPT];
#pragma unroll
// Note(Byron): upcast logits to float32
for (int ii = 0; ii < VPT; ++ii) {
row_chunk[ii] = convert_to_float<T>(row_chunk_temp[ii]);
}
// Apply tanh softcapping and correction bias
if (moe_softcapping != 0.0f || correction_bias != nullptr) {
#pragma unroll
for (int ii = 0; ii < VPT; ++ii) {
float val = row_chunk[ii];
// Apply tanh softcapping if enabled
if (moe_softcapping != 0.0f) {
val = tanhf(val / moe_softcapping) * moe_softcapping;
}
// Apply correction bias if provided
if (correction_bias != nullptr) {
/*
LDG is interleaved
|thread0 LDG| |thread1 LDG| |thread0 LDG| |thread1 LDG|
|--------- group0 --------| |----------group1 --------|
^ local2
*/
const int group_id = ii / kEltsPerLdg;
const int local_id = ii % kEltsPerLdg;
const int expert_idx = first_elt_read_by_thread +
group_id * kThreadsPerRow * kEltsPerLdg +
local_id;
val = val + correction_bias[expert_idx];
}
row_chunk[ii] = val;
}
}
// First, we perform a max reduce within the thread. We can do the max in fp16
// safely (I think) and just convert to float afterwards for the exp + sum
// reduction.
float thread_max = row_chunk[0];
#pragma unroll
for (int ii = 1; ii < VPT; ++ii) {
thread_max = max(thread_max, row_chunk[ii]);
}
/*********************************/
/********* Softmax Begin *********/
/*********************************/
// Now, we find the max within the thread group and distribute among the
// threads. We use a butterfly reduce. lane id: 0-31 within a warp
#pragma unroll
for (int mask = kThreadsPerRow / 2; mask > 0; mask /= 2) {
// butterfly reduce with (lane id ^ mask)
thread_max = max(thread_max,
XLLM_SHFL_XOR_SYNC_WIDTH(
kSoftmaxFullMask, thread_max, mask, kThreadsPerRow));
}
// From this point, thread max in all the threads have the max within the row.
// Now, we subtract the max from each element in the thread and take the exp.
// We also compute the thread local sum.
float row_sum = 0;
#pragma unroll
for (int ii = 0; ii < VPT; ++ii) {
row_chunk[ii] = expf(row_chunk[ii] - thread_max);
row_sum += row_chunk[ii];
}
// Now, we perform the sum reduce within each thread group. Similar to the max
// reduce, we use a bufferfly pattern.
#pragma unroll
for (int mask = kThreadsPerRow / 2; mask > 0; mask /= 2) {
row_sum += XLLM_SHFL_XOR_SYNC_WIDTH(
kSoftmaxFullMask, row_sum, mask, kThreadsPerRow);
}
// From this point, all threads have the max and the sum for their rows in the
// thread_max and thread_sum variables respectively. Finally, we can scale the
// rows for the softmax. Technically, for top-k gating we don't need to
// compute the entire softmax row. We can likely look at the maxes and only
// compute for the top-k values in the row. However, this kernel will likely
// not be a bottle neck and it seems better to closer match torch and find the
// argmax after computing the softmax.
const float reciprocal_row_sum = 1.f / row_sum;
#pragma unroll
for (int ii = 0; ii < VPT; ++ii) {
row_chunk[ii] = row_chunk[ii] * reciprocal_row_sum;
}
/*******************************/
/********* Softmax End *********/
/*******************************/
// Now, softmax_res contains the softmax of the row chunk. Now, I want to find
// the topk elements in each row, along with the max index.
int start_col = first_elt_read_by_thread;
static constexpr int kColsPerGroupLdg = kEltsPerLdg * kThreadsPerRow;
float row_sum_for_renormalize = 0;
for (int k_idx = 0; k_idx < k; ++k_idx) {
// First, each thread does the local argmax
float max_val = row_chunk[0];
int expert = start_col;
#pragma unroll
for (int ldg = 0, col = start_col; ldg < kLdgPerThread;
++ldg, col += kColsPerGroupLdg) {
#pragma unroll
for (int ii = 0; ii < kEltsPerLdg; ++ii) {
float val = row_chunk[ldg * kEltsPerLdg + ii];
// No check on the experts here since columns with the smallest index
// are processed first and only updated if > (not >=)
if (val > max_val) {
max_val = val;
expert = col + ii;
}
}
}
// Now, we perform the argmax reduce. We use the butterfly pattern so threads
// reach consensus about the max. This will be useful for K > 1 so that the
// threads can agree on "who" had the max value. That thread can then blank out
// their max with -inf and the warp can run more iterations...
#pragma unroll
for (int mask = kThreadsPerRow / 2; mask > 0; mask /= 2) {
float other_max = XLLM_SHFL_XOR_SYNC_WIDTH(
kSoftmaxFullMask, max_val, mask, kThreadsPerRow);
int other_expert = XLLM_SHFL_XOR_SYNC_WIDTH(
kSoftmaxFullMask, expert, mask, kThreadsPerRow);
// We want lower indices to "win" in every thread so we break ties this
// way
if (other_max > max_val ||
(other_max == max_val && other_expert < expert)) {
max_val = other_max;
expert = other_expert;
}
}
// Write the max for this k iteration to global memory.
if (thread_group_idx == 0) {
// Add a guard to ignore experts not included by this node
const bool node_uses_expert =
expert >= start_expert && expert < end_expert;
const bool should_process_row = row_is_active && node_uses_expert;
// The lead thread from each sub-group will write out the final results to
// global memory. (This will be a single) thread per row of the
// input/output matrices.
const int idx = k * thread_row + k_idx;
output[idx] = max_val;
indices[idx] = should_process_row ? (expert - start_expert) : NUM_EXPERTS;
row_sum_for_renormalize += max_val;
}
// Finally, we clear the value in the thread with the current max if there
// is another iteration to run.
if (k_idx + 1 < k) {
const int ldg_group_for_expert = expert / kColsPerGroupLdg;
const int thread_to_clear_in_group =
(expert / kEltsPerLdg) % kThreadsPerRow;
// Only the thread in the group which produced the max will reset the
// "winning" value to -inf.
if (thread_group_idx == thread_to_clear_in_group) {
const int offset_for_expert = expert % kEltsPerLdg;
// Safe to set to any negative value since row_chunk values must be
// between 0 and 1.
row_chunk[ldg_group_for_expert * kEltsPerLdg + offset_for_expert] =
-10000.f;
}
}
}
// Fuse renormalization of topk_weights into this kernel
if (renormalize && thread_group_idx == 0) {
float row_sum_for_renormalize_inv = 1.f / row_sum_for_renormalize;
#pragma unroll
for (int k_idx = 0; k_idx < k; ++k_idx) {
const int idx = k * thread_row + k_idx;
output[idx] = output[idx] * row_sum_for_renormalize_inv;
}
}
}
template <typename T, int EXPERTS, int WARPS_PER_TB>
void topk_gating_softmax_launcher_helper(const T* input,
const bool* finished,
float* output,
int* indices,
const int num_rows,
const int k,
const int start_expert,
const int end_expert,
const bool renormalize,
const float moe_softcapping,
const float* correction_bias,
cudaStream_t stream) {
static constexpr std::size_t kMaxBytesPerLdg = 16;
static constexpr int kBytesPerLdg = MIN(kMaxBytesPerLdg, sizeof(T) * EXPERTS);
using Constants = TopkConstants<T, EXPERTS, kBytesPerLdg>;
static constexpr int kVpt = Constants::VPT;
static constexpr int kRowsPerWarp = Constants::ROWS_PER_WARP;
const int num_warps = (num_rows + kRowsPerWarp - 1) / kRowsPerWarp;
const int num_blocks = (num_warps + WARPS_PER_TB - 1) / WARPS_PER_TB;
dim3 block_dim(WARP_SIZE, WARPS_PER_TB);
topk_gating_softmax<T, kVpt, EXPERTS, WARPS_PER_TB, kBytesPerLdg>
<<<num_blocks, block_dim, 0, stream>>>(input,
finished,
output,
num_rows,
indices,
k,
start_expert,
end_expert,
renormalize,
moe_softcapping,
correction_bias);
}
#define LAUNCH_SOFTMAX(TYPE, NUM_EXPERTS, WARPS_PER_TB) \
topk_gating_softmax_launcher_helper<TYPE, NUM_EXPERTS, WARPS_PER_TB>( \
gating_output, \
nullptr, \
topk_weights, \
topk_indices, \
num_tokens, \
topk, \
0, \
num_experts, \
renormalize, \
moe_softcapping, \
correction_bias, \
stream);
template <typename T>
void topk_gating_softmax_kernel_launcher(const T* gating_output,
float* topk_weights,
int* topk_indices,
float* softmax_workspace,
const int num_tokens,
const int num_experts,
const int topk,
const bool renormalize,
const float moe_softcapping,
const float* correction_bias,
cudaStream_t stream) {
static constexpr int kWarpsPerTb = 4;
switch (num_experts) {
case 1:
LAUNCH_SOFTMAX(T, 1, kWarpsPerTb);
break;
case 2:
LAUNCH_SOFTMAX(T, 2, kWarpsPerTb);
break;
case 4:
LAUNCH_SOFTMAX(T, 4, kWarpsPerTb);
break;
case 8:
LAUNCH_SOFTMAX(T, 8, kWarpsPerTb);
break;
case 16:
LAUNCH_SOFTMAX(T, 16, kWarpsPerTb);
break;
case 32:
LAUNCH_SOFTMAX(T, 32, kWarpsPerTb);
break;
case 64:
LAUNCH_SOFTMAX(T, 64, kWarpsPerTb);
break;
case 128:
LAUNCH_SOFTMAX(T, 128, kWarpsPerTb);
break;
case 256:
LAUNCH_SOFTMAX(T, 256, kWarpsPerTb);
break;
default: {
CHECK(softmax_workspace != nullptr)
<< "softmax_workspace must be provided for num_experts that are "
"not a power of 2.";
static constexpr int kTpb = 256;
moe_softmax<T, kTpb><<<num_tokens, kTpb, 0, stream>>>(gating_output,
nullptr,
softmax_workspace,
num_experts,
moe_softcapping,
correction_bias);
if (topk == 1) {
// Note: As an optimization for better performance,
// the softmax_workspace is overwritten in-place by both moeTopK and
// moe_topk_fast.
moe_topK<kTpb><<<num_tokens, kTpb, 0, stream>>>(softmax_workspace,
nullptr,
topk_weights,
topk_indices,
num_experts,
topk,
0,
num_experts,
renormalize);
} else {
moe_topk_fast<kTpb><<<num_tokens, kTpb, 0, stream>>>(softmax_workspace,
nullptr,
topk_weights,
topk_indices,
num_experts,
topk,
0,
num_experts,
renormalize);
}
}
}
}
} // namespace
namespace xllm::kernel::cuda {
void topk_softmax(torch::Tensor& topk_weights, // [num_tokens, topk]
torch::Tensor& topk_indices, // [num_tokens, topk]
torch::Tensor& gating_output, // [num_tokens, num_experts]
const bool renormalize,
const double moe_softcapping,
const std::optional<torch::Tensor>& correction_bias) {
// Check data type
CHECK(gating_output.scalar_type() == at::ScalarType::Float ||
gating_output.scalar_type() == at::ScalarType::Half ||
gating_output.scalar_type() == at::ScalarType::BFloat16)
<< "gating_output must be float32, float16, or bfloat16";
// Check dimensions
CHECK(gating_output.dim() == 2)
<< "gating_output must be 2D tensor [num_tokens, num_experts]";
CHECK(topk_weights.dim() == 2)
<< "topk_weights must be 2D tensor [num_tokens, topk]";
CHECK(topk_indices.dim() == 2)
<< "topk_indices must be 2D tensor [num_tokens, topk]";
// Check shapes
CHECK(gating_output.size(0) == topk_weights.size(0))
<< "First dimension of topk_weights must match num_tokens in "
"gating_output"
<< "First dimension of topk_indices must match num_tokens in "
"gating_output";
CHECK(topk_weights.size(-1) == topk_indices.size(-1))
<< "Second dimension of topk_indices must match topk in topk_weights"
<< "topk must be less than or equal to num_experts";
const int num_experts = static_cast<int>(gating_output.size(-1));
const int num_tokens = static_cast<int>(gating_output.size(0));
const int topk = static_cast<int>(topk_weights.size(-1));
const bool is_pow_2 =
(num_experts != 0) && ((num_experts & (num_experts - 1)) == 0);
const bool needs_workspace = !is_pow_2 || num_experts > 256;
const int64_t workspace_size = needs_workspace ? num_tokens * num_experts : 0;
const at::cuda::OptionalCUDAGuard device_guard(device_of(gating_output));
const cudaStream_t stream = at::cuda::getCurrentCUDAStream();
torch::Tensor softmax_workspace = torch::empty(
{workspace_size}, gating_output.options().dtype(at::ScalarType::Float));
const at::ScalarType dtype = gating_output.scalar_type();
// Validate correction_bias if provided - must always be float32
const float* bias_ptr = nullptr;
if (correction_bias.has_value()) {
const torch::Tensor& bias_tensor = correction_bias.value();
CHECK(bias_tensor.dim() == 1)
<< "correction_bias must be 1D tensor [num_experts]";
CHECK(bias_tensor.size(0) == num_experts)
<< "correction_bias size must match num_experts";
CHECK(bias_tensor.scalar_type() == at::ScalarType::Float)
<< "correction_bias must be float32, got " << bias_tensor.scalar_type();
bias_ptr = bias_tensor.data_ptr<float>();
}
// Cast moe_softcapping from double to float for CUDA kernels
const float moe_softcapping_f = static_cast<float>(moe_softcapping);
if (dtype == at::ScalarType::Float) {
topk_gating_softmax_kernel_launcher<float>(
gating_output.data_ptr<float>(),
topk_weights.data_ptr<float>(),
topk_indices.data_ptr<int>(),
softmax_workspace.data_ptr<float>(),
num_tokens,
num_experts,
topk,
renormalize,
moe_softcapping_f,
bias_ptr,
stream);
} else if (dtype == at::ScalarType::Half) {
topk_gating_softmax_kernel_launcher<__half>(
reinterpret_cast<const __half*>(gating_output.data_ptr<at::Half>()),
topk_weights.data_ptr<float>(),
topk_indices.data_ptr<int>(),
softmax_workspace.data_ptr<float>(),
num_tokens,
num_experts,
topk,
renormalize,
moe_softcapping_f,
bias_ptr,
stream);
} else if (dtype == at::ScalarType::BFloat16) {
topk_gating_softmax_kernel_launcher<BFloat16Type>(
reinterpret_cast<const BFloat16Type*>(
gating_output.data_ptr<at::BFloat16>()),
topk_weights.data_ptr<float>(),
topk_indices.data_ptr<int>(),
softmax_workspace.data_ptr<float>(),
num_tokens,
num_experts,
topk,
renormalize,
moe_softcapping_f,
bias_ptr,
stream);
} else {
LOG(FATAL) << "Unsupported gating_output dtype: " << dtype;
}
}
} // namespace xllm::kernel::cuda

1164
upstream_ref/xllm_latest/core/layers/npu_torch/qwen3_gated_delta_net_base.cpp Normal file
View File

File diff suppressed because it is too large Load Diff

112
upstream_ref/xllm_latest/core/layers/npu_torch/qwen3_gated_delta_net_base.h Normal file
View File

@@ -0,0 +1,112 @@
/* Copyright 2025-2026 The xLLM Authors.
Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at
https://github.com/jd-opensource/xllm/blob/main/LICENSE
Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
==============================================================================*/
#pragma once
#include <torch/torch.h>
#include <optional>
#include <string>
#include <tuple>
#include <utility>
#include "attention.h"
#include "framework/kv_cache/kv_cache.h"
#include "framework/model/model_args.h"
#include "framework/parallel_state/parallel_args.h"
#include "framework/quant_args.h"
#include "framework/state_dict/state_dict.h"
#include "framework/state_dict/utils.h"
#include "layers/common/linear.h"
#include "layers/common/rms_norm_gated.h"
namespace xllm {
namespace layer {
class Qwen3GatedDeltaNetBaseImpl : public torch::nn::Module {
public:
Qwen3GatedDeltaNetBaseImpl() = default;
Qwen3GatedDeltaNetBaseImpl(const ModelArgs& args,
const QuantArgs& quant_args,
const ParallelArgs& parallel_args,
const torch::TensorOptions& options);
virtual void load_state_dict(const StateDict& state_dict) = 0;
virtual void verify_loaded_weights(const std::string& prefix) const = 0;
torch::Tensor forward(const torch::Tensor& hidden_states,
const AttentionMetadata& attn_metadata,
KVCache& kv_cache,
const ModelInputParams& input_params);
protected:
virtual std::pair<torch::Tensor, torch::Tensor> project_decode_inputs(
const torch::Tensor& hidden_states) = 0;
virtual std::pair<torch::Tensor, torch::Tensor> project_flat_inputs(
const torch::Tensor& hidden_states) = 0;
// Qwen3.5 overrides this to project and reshape its separate qkv/z/b/a
// weights in every forward mode. Qwen3Next keeps qkvz/ba packed and returns
// nullopt to select the fused-split fallback.
virtual std::optional<
std::tuple<torch::Tensor, torch::Tensor, torch::Tensor, torch::Tensor>>
project_split_inputs(const torch::Tensor& hidden_states,
const AttentionMetadata& attn_metadata) {
return std::nullopt;
}
virtual bool use_fla_ssm_state_layout() const { return false; }
void load_common_state_dict(const StateDict& state_dict);
void verify_common_loaded_weights(const std::string& prefix) const;
torch::Tensor get_linear_state_indices(const ModelInputParams& input_params,
const torch::Device& device) const;
std::pair<torch::Tensor, torch::Tensor> project_padded_inputs(
const torch::Tensor& hidden_states,
const AttentionMetadata& attn_metadata);
torch::Tensor reshape_qkvz_unpad(const AttentionMetadata& attn_metadata,
const torch::Tensor& padded_qkvz) const;
// Projection outputs are packed as [total_tokens, dim], while GDN kernels
// consume dense [batch, max_query_len, dim] tensors. Split the packed tokens
// by query length and pad each sequence before entering the kernels.
torch::Tensor reshape_projected_tokens_with_pad(
const AttentionMetadata& attn_metadata,
const torch::Tensor& projected_tokens) const;
std::tuple<torch::Tensor, torch::Tensor, torch::Tensor> process_mixed_qkv(
torch::Tensor& mixed_qkv) const;
int64_t num_k_heads_ = 0;
int64_t num_v_heads_ = 0;
int64_t head_k_dim_ = 0;
int64_t head_v_dim_ = 0;
int64_t k_size_ = 0;
int64_t v_size_ = 0;
int64_t tp_size_ = 1;
int64_t rank_ = 0;
int32_t conv_kernel_size_ = 0;
ColumnParallelLinear conv1d_{nullptr};
RowParallelLinear o_proj_{nullptr};
RmsNormGated norm_{nullptr};
DEFINE_WEIGHT(dt_bias);
DEFINE_WEIGHT(A_log);
};
} // namespace layer
} // namespace xllm