0359103b9b
cuda_ops_api.h includes glog/logging.h and ATen/DynamicLibrary.h which are not available in corex standalone compilation. All kernel .cu files only need device_utils.cuh (provides namespace, XLLM_KERNEL_ATTR macro, CUB includes, type helpers). Fixed files: norm.cu, activation.cu, rope.cu, block_copy.cu, reshape_paged_cache.cu moe/moe_combine.cu, moe/moe_compute_index.cu, moe/moe_fused_topk.cu
106 lines
4.0 KiB
Plaintext
Executable File
106 lines
4.0 KiB
Plaintext
Executable File
/* 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 "device_utils.cuh"
|
|
|
|
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
|