fix(critical): fold max_completion_tokens + max_num_seqs=2 + max_model_len=80000 + xllm_latest layer import
Sub 655 root causes (confirmed from log analysis): 1. protocol.py: max_completion_tokens never folded into max_tokens → 162/881 replay requests rejected 400 (extra_forbidden) 2. max_num_seqs=1 → t2_n_2 test fails (needs n=2) 3. max_model_len=131072 → OOM crash at 62% replay, opencompass all 0 Fixes: - protocol.py: model_validator fold_max_completion_tokens - yaml: max_num_seqs=2, max_model_len=80000, PYTORCH_CUDA_ALLOC_CONF - topk_softmax stays =0 (corex CUB BlockReduce incompatible on BI-V100) xllm_latest import to ex_engine/: - npu_torch layers: GDN(1164L), Qwen3.5 GDN, attention, fused_moe - cuda/moe kernels: topk_softmax_kernels.cuh, moe_combine, moe_compute_index - npu kernels: causal_conv1d, recurrent_gated_delta_rule - model headers: qwen3_5.h, qwen3_next.h
This commit is contained in:
project6-dev
124
ex_engine/xllm_kernels/cuda/moe/fused_moe.cpp
Normal file
124
ex_engine/xllm_kernels/cuda/moe/fused_moe.cpp
Normal 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
ex_engine/xllm_kernels/cuda/moe/moe_combine.cu
Executable file
105
ex_engine/xllm_kernels/cuda/moe/moe_combine.cu
Executable 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
ex_engine/xllm_kernels/cuda/moe/moe_compute_index.cu
Normal file
155
ex_engine/xllm_kernels/cuda/moe/moe_compute_index.cu
Normal 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
ex_engine/xllm_kernels/cuda/moe/moe_fused_topk.cu
Normal file
59
ex_engine/xllm_kernels/cuda/moe/moe_fused_topk.cu
Normal 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
ex_engine/xllm_kernels/cuda/moe/moe_topk.cuh
Normal file
345
ex_engine/xllm_kernels/cuda/moe/moe_topk.cuh
Normal 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
ex_engine/xllm_kernels/cuda/moe/moe_topk_sigmoid_kernels.cuh
Normal file
609
ex_engine/xllm_kernels/cuda/moe/moe_topk_sigmoid_kernels.cuh
Normal 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
ex_engine/xllm_kernels/cuda/moe/moe_topk_softmax_kernels.cuh
Normal file
867
ex_engine/xllm_kernels/cuda/moe/moe_topk_softmax_kernels.cuh
Normal 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
|
||||
59
ex_engine/xllm_kernels/npu/npu_causal_conv1d.cpp
Normal file
59
ex_engine/xllm_kernels/npu/npu_causal_conv1d.cpp
Normal file
@@ -0,0 +1,59 @@
|
||||
/* Copyright 2026 The xLLM Authors. All Rights Reserved.
|
||||
|
||||
Licensed under the Apache License, Version 2.0 (the "License");
|
||||
you may not use this file except in compliance with the License.
|
||||
You may obtain a copy of the License at
|
||||
|
||||
https://github.com/jd-opensource/xllm/blob/main/LICENSE
|
||||
|
||||
Unless required by applicable law or agreed to in writing, software
|
||||
distributed under the License is distributed on an "AS IS" BASIS,
|
||||
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
See the License for the specific language governing permissions and
|
||||
limitations under the License.
|
||||
==============================================================================*/
|
||||
|
||||
#include "core/kernels/npu/aclnn/pytorch_npu_helper.hpp"
|
||||
#include "core/kernels/npu/utils.h"
|
||||
#include "core/kernels/npu/xllm_ops/xllm_ops_api.h"
|
||||
|
||||
namespace xllm::kernel::npu {
|
||||
|
||||
torch::Tensor causal_conv1d(const torch::Tensor& x,
|
||||
const torch::Tensor& weight,
|
||||
const torch::Tensor& conv_state,
|
||||
const std::optional<torch::Tensor>& bias_opt,
|
||||
const torch::IntArrayRef query_start_loc_opt,
|
||||
const torch::IntArrayRef cache_indices_opt,
|
||||
const torch::IntArrayRef initial_state_mode_opt,
|
||||
const torch::IntArrayRef num_accepted_tokens_opt,
|
||||
int64_t activation_mode,
|
||||
int64_t pad_slot_id,
|
||||
int64_t run_mode) {
|
||||
check_tensor(x, "x", "causal_conv1d");
|
||||
check_tensor(weight, "weight", "causal_conv1d");
|
||||
check_tensor(conv_state, "conv_state", "causal_conv1d");
|
||||
|
||||
c10::optional<torch::Tensor> bias_tensor = c10::nullopt;
|
||||
if (bias_opt.has_value() && bias_opt.value().defined()) {
|
||||
bias_tensor = bias_opt.value();
|
||||
}
|
||||
|
||||
torch::Tensor output = torch::empty(x.sizes(), x.options());
|
||||
EXEC_NPU_CMD(aclnnCausalConv1d,
|
||||
x,
|
||||
weight,
|
||||
bias_tensor,
|
||||
conv_state,
|
||||
query_start_loc_opt,
|
||||
cache_indices_opt,
|
||||
initial_state_mode_opt,
|
||||
num_accepted_tokens_opt,
|
||||
activation_mode,
|
||||
pad_slot_id,
|
||||
run_mode,
|
||||
output);
|
||||
return output;
|
||||
}
|
||||
|
||||
} // namespace xllm::kernel::npu
|
||||
@@ -0,0 +1,83 @@
|
||||
/* Copyright 2026 The xLLM Authors. All Rights Reserved.
|
||||
|
||||
Licensed under the Apache License, Version 2.0 (the "License");
|
||||
you may not use this file except in compliance with the License.
|
||||
You may obtain a copy of the License at
|
||||
|
||||
https://github.com/jd-opensource/xllm/blob/main/LICENSE
|
||||
|
||||
Unless required by applicable law or agreed to in writing, software
|
||||
distributed under the License is distributed on an "AS IS" BASIS,
|
||||
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
See the License for the specific language governing permissions and
|
||||
limitations under the License.
|
||||
==============================================================================*/
|
||||
|
||||
#include <glog/logging.h>
|
||||
|
||||
#include "core/kernels/npu/aclnn/pytorch_npu_helper.hpp"
|
||||
#include "core/kernels/npu/npu_ops_api.h"
|
||||
#include "core/kernels/npu/utils.h"
|
||||
|
||||
namespace {
|
||||
|
||||
c10::optional<torch::Tensor> to_c10_optional_tensor(
|
||||
const std::optional<torch::Tensor>& tensor_opt) {
|
||||
if (tensor_opt.has_value() && tensor_opt.value().defined()) {
|
||||
return tensor_opt.value();
|
||||
}
|
||||
return c10::nullopt;
|
||||
}
|
||||
|
||||
} // namespace
|
||||
|
||||
namespace xllm::kernel::npu {
|
||||
|
||||
torch::Tensor npu_recurrent_gated_delta_rule(
|
||||
const torch::Tensor& query,
|
||||
const torch::Tensor& key,
|
||||
const torch::Tensor& value,
|
||||
torch::Tensor& state,
|
||||
const std::optional<torch::Tensor>& beta,
|
||||
const std::optional<double> scale,
|
||||
const std::optional<torch::Tensor>& actual_seq_lengths,
|
||||
const std::optional<torch::Tensor>& ssm_state_indices,
|
||||
const std::optional<torch::Tensor>& num_accepted_tokens,
|
||||
const std::optional<torch::Tensor>& g,
|
||||
const std::optional<torch::Tensor>& gk) {
|
||||
check_tensor(query, "query", "recurrent_gated_delta_rule");
|
||||
check_tensor(key, "key", "recurrent_gated_delta_rule");
|
||||
check_tensor(value, "value", "recurrent_gated_delta_rule");
|
||||
check_tensor(state, "state", "recurrent_gated_delta_rule");
|
||||
CHECK(scale.has_value())
|
||||
<< "recurrent_gated_delta_rule requires a valid scale value";
|
||||
|
||||
c10::optional<torch::Tensor> beta_tensor = to_c10_optional_tensor(beta);
|
||||
c10::optional<torch::Tensor> actual_seq_lengths_tensor =
|
||||
to_c10_optional_tensor(actual_seq_lengths);
|
||||
c10::optional<torch::Tensor> ssm_state_indices_tensor =
|
||||
to_c10_optional_tensor(ssm_state_indices);
|
||||
c10::optional<torch::Tensor> num_accepted_tokens_tensor =
|
||||
to_c10_optional_tensor(num_accepted_tokens);
|
||||
c10::optional<torch::Tensor> g_tensor = to_c10_optional_tensor(g);
|
||||
c10::optional<torch::Tensor> gk_tensor = to_c10_optional_tensor(gk);
|
||||
float scale_value = static_cast<float>(scale.value());
|
||||
torch::Tensor output = torch::empty_like(value);
|
||||
|
||||
EXEC_NPU_CMD(aclnnRecurrentGatedDeltaRule,
|
||||
query,
|
||||
key,
|
||||
value,
|
||||
beta_tensor,
|
||||
state,
|
||||
actual_seq_lengths_tensor,
|
||||
ssm_state_indices_tensor,
|
||||
g_tensor,
|
||||
gk_tensor,
|
||||
num_accepted_tokens_tensor,
|
||||
scale_value,
|
||||
output);
|
||||
return output;
|
||||
}
|
||||
|
||||
} // namespace xllm::kernel::npu
|
||||
Reference in New Issue
Block a user