355 lines
12 KiB
Plaintext
355 lines
12 KiB
Plaintext
#include <ATen/cuda/CUDAContext.h>
|
|
#include <cuda_runtime.h>
|
|
#include <cutlass/array.h>
|
|
#include <cutlass/cutlass.h>
|
|
#include <cutlass/numeric_types.h>
|
|
#include <torch/all.h>
|
|
|
|
#include <cfloat>
|
|
|
|
using bfloat16_t = cutlass::bfloat16_t;
|
|
using float16_t = cutlass::half_t;
|
|
|
|
// Kimi K2 specific constants
|
|
static constexpr int WARP_SIZE = 32;
|
|
static constexpr int WARPS_PER_CTA = 6;
|
|
static constexpr int NUM_EXPERTS = 384;
|
|
static constexpr int VPT = 12; // 384 / 32 = 12
|
|
|
|
// Small token optimization constants
|
|
static constexpr int SMALL_TOKEN_THRESHOLD = 512;
|
|
static constexpr int WARPS_PER_TOKEN_SMALL = 12; // Use 12 warps per token for small batches
|
|
static constexpr int THREADS_PER_BLOCK_SMALL = WARPS_PER_TOKEN_SMALL * WARP_SIZE; // 384 threads
|
|
|
|
template <typename T>
|
|
__device__ inline bool cmp_gt(const T& a, const T& b) {
|
|
return static_cast<float>(a) > static_cast<float>(b);
|
|
}
|
|
|
|
template <typename T>
|
|
__device__ inline bool cmp_eq(const T& a, const T& b) {
|
|
return static_cast<float>(a) == static_cast<float>(b);
|
|
}
|
|
|
|
// Small token optimized kernel: Multiple warps collaborate on a single token
|
|
template <typename T>
|
|
__global__ void kimi_k2_moe_fused_gate_kernel_small_token(
|
|
T* input,
|
|
T* bias,
|
|
float* output_ptr,
|
|
int32_t* indices_ptr,
|
|
int64_t num_rows,
|
|
int64_t topk,
|
|
bool renormalize,
|
|
double routed_scaling_factor,
|
|
bool apply_routed_scaling_factor_on_output) {
|
|
// Each block handles one token with WARPS_PER_TOKEN_SMALL warps collaborating
|
|
int64_t row_idx = blockIdx.x;
|
|
if (row_idx >= num_rows) return;
|
|
|
|
int tid = threadIdx.x;
|
|
int warp_id = tid / WARP_SIZE;
|
|
int lane_id = tid % WARP_SIZE;
|
|
|
|
// Shared memory for all warps to collaborate
|
|
__shared__ float shared_scores[NUM_EXPERTS];
|
|
__shared__ float shared_original_scores[NUM_EXPERTS];
|
|
|
|
// Each thread loads one expert (384 threads for 384 experts)
|
|
if (tid < NUM_EXPERTS) {
|
|
T input_val = input[row_idx * NUM_EXPERTS + tid];
|
|
T bias_val = bias[tid];
|
|
float sigmoid_val = 1.0f / (1.0f + expf(-static_cast<float>(input_val)));
|
|
float biased_val = sigmoid_val + static_cast<float>(bias_val);
|
|
shared_scores[tid] = biased_val;
|
|
shared_original_scores[tid] = sigmoid_val;
|
|
}
|
|
|
|
__syncthreads();
|
|
|
|
// Parallel TopK: Each warp processes a portion of experts
|
|
// Use multiple warps to find top-k elements in parallel
|
|
int experts_per_warp = (NUM_EXPERTS + WARPS_PER_TOKEN_SMALL - 1) / WARPS_PER_TOKEN_SMALL;
|
|
int warp_start = warp_id * experts_per_warp;
|
|
int warp_end = min(warp_start + experts_per_warp, NUM_EXPERTS);
|
|
|
|
for (int k = 0; k < topk; k++) {
|
|
float max_val = -FLT_MAX;
|
|
int max_expert = -1;
|
|
|
|
// Each warp finds the max in its portion
|
|
for (int expert = warp_start + lane_id; expert < warp_end; expert += WARP_SIZE) {
|
|
float val = shared_scores[expert];
|
|
if (val > max_val) {
|
|
max_val = val;
|
|
max_expert = expert;
|
|
}
|
|
}
|
|
|
|
// Warp-level reduction to find warp's maximum
|
|
for (int offset = WARP_SIZE / 2; offset > 0; offset /= 2) {
|
|
float other_val = __shfl_down_sync(0xFFFFFFFF, max_val, offset);
|
|
int other_expert = __shfl_down_sync(0xFFFFFFFF, max_expert, offset);
|
|
|
|
if (other_val > max_val || (other_val == max_val && other_expert < max_expert)) {
|
|
max_val = other_val;
|
|
max_expert = other_expert;
|
|
}
|
|
}
|
|
|
|
// Store warp results in shared memory
|
|
__shared__ float warp_max_vals[WARPS_PER_TOKEN_SMALL];
|
|
__shared__ int warp_max_experts[WARPS_PER_TOKEN_SMALL];
|
|
|
|
if (lane_id == 0) {
|
|
warp_max_vals[warp_id] = max_val;
|
|
warp_max_experts[warp_id] = max_expert;
|
|
}
|
|
|
|
__syncthreads();
|
|
|
|
// First warp reduces across all warp results
|
|
if (warp_id == 0) {
|
|
float final_max_val = -FLT_MAX;
|
|
int final_max_expert = -1;
|
|
|
|
if (lane_id < WARPS_PER_TOKEN_SMALL) {
|
|
final_max_val = warp_max_vals[lane_id];
|
|
final_max_expert = warp_max_experts[lane_id];
|
|
}
|
|
|
|
// Warp reduction
|
|
for (int offset = WARP_SIZE / 2; offset > 0; offset /= 2) {
|
|
float other_val = __shfl_down_sync(0xFFFFFFFF, final_max_val, offset);
|
|
int other_expert = __shfl_down_sync(0xFFFFFFFF, final_max_expert, offset);
|
|
|
|
if (other_val > final_max_val || (other_val == final_max_val && other_expert < final_max_expert)) {
|
|
final_max_val = other_val;
|
|
final_max_expert = other_expert;
|
|
}
|
|
}
|
|
|
|
// Lane 0 writes result and marks the expert as used
|
|
if (lane_id == 0 && final_max_expert != -1) {
|
|
int64_t output_idx = row_idx * topk + k;
|
|
output_ptr[output_idx] = shared_original_scores[final_max_expert];
|
|
indices_ptr[output_idx] = final_max_expert;
|
|
shared_scores[final_max_expert] = -FLT_MAX;
|
|
}
|
|
}
|
|
|
|
__syncthreads();
|
|
}
|
|
|
|
// Renormalization (only first warp)
|
|
if (renormalize && warp_id == 0 && lane_id == 0) {
|
|
float sum = 0.0f;
|
|
for (int k = 0; k < topk; k++) {
|
|
sum += output_ptr[row_idx * topk + k];
|
|
}
|
|
|
|
if (sum > 0.0f) {
|
|
for (int k = 0; k < topk; k++) {
|
|
int64_t idx = row_idx * topk + k;
|
|
output_ptr[idx] /= sum;
|
|
if (apply_routed_scaling_factor_on_output) {
|
|
output_ptr[idx] *= static_cast<float>(routed_scaling_factor);
|
|
}
|
|
}
|
|
}
|
|
}
|
|
}
|
|
|
|
template <typename T>
|
|
__global__ void kimi_k2_moe_fused_gate_kernel(
|
|
T* input,
|
|
T* bias,
|
|
float* output_ptr,
|
|
int32_t* indices_ptr,
|
|
int64_t num_rows,
|
|
int64_t topk,
|
|
bool renormalize,
|
|
double routed_scaling_factor,
|
|
bool apply_routed_scaling_factor_on_output) {
|
|
int64_t row_idx = blockIdx.x * WARPS_PER_CTA + threadIdx.y;
|
|
if (row_idx >= num_rows) return;
|
|
|
|
int lane_id = threadIdx.x;
|
|
int warp_id = threadIdx.y;
|
|
|
|
__shared__ float shared_scores[NUM_EXPERTS * WARPS_PER_CTA];
|
|
__shared__ float shared_original_scores[NUM_EXPERTS * WARPS_PER_CTA];
|
|
|
|
float* warp_scores = shared_scores + warp_id * NUM_EXPERTS;
|
|
float* warp_original_scores = shared_original_scores + warp_id * NUM_EXPERTS;
|
|
|
|
for (int expert = lane_id; expert < NUM_EXPERTS; expert += WARP_SIZE) {
|
|
T input_val = input[row_idx * NUM_EXPERTS + expert];
|
|
T bias_val = bias[expert];
|
|
float sigmoid_val = 1.0f / (1.0f + expf(-static_cast<float>(input_val)));
|
|
float biased_val = sigmoid_val + static_cast<float>(bias_val);
|
|
warp_scores[expert] = biased_val;
|
|
warp_original_scores[expert] = sigmoid_val;
|
|
}
|
|
|
|
__syncthreads();
|
|
|
|
for (int k = 0; k < topk; k++) {
|
|
float max_val = -FLT_MAX;
|
|
int max_expert = -1;
|
|
|
|
for (int expert = lane_id; expert < NUM_EXPERTS; expert += WARP_SIZE) {
|
|
if (warp_scores[expert] > max_val) {
|
|
max_val = warp_scores[expert];
|
|
max_expert = expert;
|
|
}
|
|
}
|
|
|
|
for (int offset = WARP_SIZE / 2; offset > 0; offset /= 2) {
|
|
float other_val = __shfl_down_sync(0xFFFFFFFF, max_val, offset);
|
|
int other_expert = __shfl_down_sync(0xFFFFFFFF, max_expert, offset);
|
|
|
|
if (other_val > max_val || (other_val == max_val && other_expert < max_expert)) {
|
|
max_val = other_val;
|
|
max_expert = other_expert;
|
|
}
|
|
}
|
|
|
|
if (lane_id == 0 && max_expert != -1) {
|
|
int64_t output_idx = row_idx * topk + k;
|
|
output_ptr[output_idx] = warp_original_scores[max_expert];
|
|
indices_ptr[output_idx] = max_expert;
|
|
warp_scores[max_expert] = -FLT_MAX;
|
|
}
|
|
|
|
__syncwarp();
|
|
}
|
|
|
|
__syncthreads();
|
|
|
|
if (renormalize && lane_id == 0) {
|
|
float sum = 0.0f;
|
|
for (int k = 0; k < topk; k++) {
|
|
sum += output_ptr[row_idx * topk + k];
|
|
}
|
|
|
|
if (sum > 0.0f) {
|
|
for (int k = 0; k < topk; k++) {
|
|
int64_t idx = row_idx * topk + k;
|
|
output_ptr[idx] /= sum;
|
|
if (apply_routed_scaling_factor_on_output) {
|
|
output_ptr[idx] *= static_cast<float>(routed_scaling_factor);
|
|
}
|
|
}
|
|
}
|
|
}
|
|
}
|
|
|
|
std::vector<at::Tensor> kimi_k2_moe_fused_gate(
|
|
at::Tensor& input,
|
|
at::Tensor& bias,
|
|
int64_t topk,
|
|
bool renormalize,
|
|
double routed_scaling_factor,
|
|
bool apply_routed_scaling_factor_on_output) {
|
|
int64_t num_rows = input.size(0);
|
|
int32_t num_experts = input.size(1);
|
|
|
|
// Assert: Only support 384 experts
|
|
TORCH_CHECK(num_experts == 384, "kimi_k2_moe_fused_gate only supports 384 experts, but got ", num_experts);
|
|
TORCH_CHECK(input.dtype() == bias.dtype(), "input and bias should have the same dtype");
|
|
|
|
auto options = torch::TensorOptions().dtype(torch::kFloat32).device(torch::kCUDA);
|
|
auto output = torch::empty({num_rows, topk}, options);
|
|
auto indices = torch::empty({num_rows, topk}, options.dtype(torch::kInt32));
|
|
|
|
const cudaStream_t stream = at::cuda::getCurrentCUDAStream();
|
|
|
|
bool use_small_token_kernel = num_rows <= SMALL_TOKEN_THRESHOLD;
|
|
|
|
if (use_small_token_kernel) {
|
|
// Small token kernel: Each block handles 1 token with multiple warps collaborating
|
|
int64_t num_blocks = num_rows;
|
|
dim3 block_dim(THREADS_PER_BLOCK_SMALL);
|
|
|
|
if (input.scalar_type() == at::kBFloat16) {
|
|
kimi_k2_moe_fused_gate_kernel_small_token<bfloat16_t><<<num_blocks, block_dim, 0, stream>>>(
|
|
reinterpret_cast<bfloat16_t*>(input.data_ptr()),
|
|
reinterpret_cast<bfloat16_t*>(bias.data_ptr()),
|
|
output.data_ptr<float>(),
|
|
indices.data_ptr<int32_t>(),
|
|
num_rows,
|
|
topk,
|
|
renormalize,
|
|
routed_scaling_factor,
|
|
apply_routed_scaling_factor_on_output);
|
|
} else if (input.scalar_type() == at::kHalf) {
|
|
kimi_k2_moe_fused_gate_kernel_small_token<float16_t><<<num_blocks, block_dim, 0, stream>>>(
|
|
reinterpret_cast<float16_t*>(input.data_ptr()),
|
|
reinterpret_cast<float16_t*>(bias.data_ptr()),
|
|
output.data_ptr<float>(),
|
|
indices.data_ptr<int32_t>(),
|
|
num_rows,
|
|
topk,
|
|
renormalize,
|
|
routed_scaling_factor,
|
|
apply_routed_scaling_factor_on_output);
|
|
} else if (input.scalar_type() == at::kFloat) {
|
|
kimi_k2_moe_fused_gate_kernel_small_token<float><<<num_blocks, block_dim, 0, stream>>>(
|
|
input.data_ptr<float>(),
|
|
bias.data_ptr<float>(),
|
|
output.data_ptr<float>(),
|
|
indices.data_ptr<int32_t>(),
|
|
num_rows,
|
|
topk,
|
|
renormalize,
|
|
routed_scaling_factor,
|
|
apply_routed_scaling_factor_on_output);
|
|
} else {
|
|
TORCH_CHECK(false, "Unsupported data type for kimi_k2_moe_fused_gate");
|
|
}
|
|
} else {
|
|
int64_t num_blocks = (num_rows + WARPS_PER_CTA - 1) / WARPS_PER_CTA;
|
|
dim3 block_dim(WARP_SIZE, WARPS_PER_CTA);
|
|
|
|
if (input.scalar_type() == at::kBFloat16) {
|
|
kimi_k2_moe_fused_gate_kernel<bfloat16_t><<<num_blocks, block_dim, 0, stream>>>(
|
|
reinterpret_cast<bfloat16_t*>(input.data_ptr()),
|
|
reinterpret_cast<bfloat16_t*>(bias.data_ptr()),
|
|
output.data_ptr<float>(),
|
|
indices.data_ptr<int32_t>(),
|
|
num_rows,
|
|
topk,
|
|
renormalize,
|
|
routed_scaling_factor,
|
|
apply_routed_scaling_factor_on_output);
|
|
} else if (input.scalar_type() == at::kHalf) {
|
|
kimi_k2_moe_fused_gate_kernel<float16_t><<<num_blocks, block_dim, 0, stream>>>(
|
|
reinterpret_cast<float16_t*>(input.data_ptr()),
|
|
reinterpret_cast<float16_t*>(bias.data_ptr()),
|
|
output.data_ptr<float>(),
|
|
indices.data_ptr<int32_t>(),
|
|
num_rows,
|
|
topk,
|
|
renormalize,
|
|
routed_scaling_factor,
|
|
apply_routed_scaling_factor_on_output);
|
|
} else if (input.scalar_type() == at::kFloat) {
|
|
kimi_k2_moe_fused_gate_kernel<float><<<num_blocks, block_dim, 0, stream>>>(
|
|
input.data_ptr<float>(),
|
|
bias.data_ptr<float>(),
|
|
output.data_ptr<float>(),
|
|
indices.data_ptr<int32_t>(),
|
|
num_rows,
|
|
topk,
|
|
renormalize,
|
|
routed_scaling_factor,
|
|
apply_routed_scaling_factor_on_output);
|
|
} else {
|
|
TORCH_CHECK(false, "Unsupported data type for kimi_k2_moe_fused_gate");
|
|
}
|
|
}
|
|
|
|
return {output, indices};
|
|
}
|