0359103b9b
cuda_ops_api.h includes glog/logging.h and ATen/DynamicLibrary.h which are not available in corex standalone compilation. All kernel .cu files only need device_utils.cuh (provides namespace, XLLM_KERNEL_ATTR macro, CUB includes, type helpers). Fixed files: norm.cu, activation.cu, rope.cu, block_copy.cu, reshape_paged_cache.cu moe/moe_combine.cu, moe/moe_compute_index.cu, moe/moe_fused_topk.cu
259 lines
10 KiB
Plaintext
259 lines
10 KiB
Plaintext
/* Copyright 2025 The vLLM Authors and 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 <ATen/cuda/CUDAContext.h>
|
|
#include <c10/cuda/CUDAGuard.h>
|
|
#include <torch/all.h>
|
|
|
|
|
|
#include "device_utils.cuh"
|
|
|
|
// ref to:
|
|
// https://github.com/vllm-project/vllm/blob/main/csrc/pos_encoding_kernels.cu
|
|
|
|
namespace {
|
|
|
|
template <typename scalar_t, bool IS_NEOX>
|
|
inline __device__ void apply_token_rotary_embedding(
|
|
scalar_t* __restrict__ arr,
|
|
const scalar_t* __restrict__ cos_ptr,
|
|
const scalar_t* __restrict__ sin_ptr,
|
|
int rot_offset,
|
|
int embed_dim) {
|
|
int x_index, y_index;
|
|
scalar_t cos, sin;
|
|
if (IS_NEOX) {
|
|
// GPT-NeoX style rotary embedding.
|
|
x_index = rot_offset;
|
|
y_index = embed_dim + rot_offset;
|
|
cos = *(cos_ptr + x_index);
|
|
sin = *(sin_ptr + x_index);
|
|
} else {
|
|
// GPT-J style rotary embedding.
|
|
x_index = 2 * rot_offset;
|
|
y_index = 2 * rot_offset + 1;
|
|
cos = *(cos_ptr + x_index / 2);
|
|
sin = *(sin_ptr + x_index / 2);
|
|
}
|
|
|
|
const scalar_t x = arr[x_index];
|
|
const scalar_t y = arr[y_index];
|
|
arr[x_index] = x * cos - y * sin;
|
|
arr[y_index] = y * cos + x * sin;
|
|
}
|
|
|
|
template <typename scalar_t, bool IS_NEOX>
|
|
inline __device__ void apply_rotary_embedding(
|
|
scalar_t* __restrict__ query, // [batch_size, seq_len, num_heads,
|
|
// head_size] or [num_tokens, num_heads,
|
|
// head_size]
|
|
scalar_t* __restrict__ key, // nullptr or
|
|
// [batch_size, seq_len, num_kv_heads,
|
|
// head_size] or [num_tokens, num_kv_heads,
|
|
// head_size]
|
|
const scalar_t* cache_ptr,
|
|
const int head_size,
|
|
const int num_heads,
|
|
const int num_kv_heads,
|
|
const int rot_dim,
|
|
const int token_idx,
|
|
const int64_t query_stride,
|
|
const int64_t key_stride,
|
|
const int64_t head_stride) {
|
|
const int embed_dim = rot_dim / 2;
|
|
const scalar_t* cos_ptr = cache_ptr;
|
|
const scalar_t* sin_ptr = cache_ptr + embed_dim;
|
|
|
|
const int nq = num_heads * embed_dim;
|
|
for (int i = threadIdx.x; i < nq; i += blockDim.x) {
|
|
const int head_idx = i / embed_dim;
|
|
const int64_t token_head =
|
|
token_idx * query_stride + head_idx * head_stride;
|
|
const int rot_offset = i % embed_dim;
|
|
apply_token_rotary_embedding<scalar_t, IS_NEOX>(
|
|
query + token_head, cos_ptr, sin_ptr, rot_offset, embed_dim);
|
|
}
|
|
|
|
if (key != nullptr) {
|
|
const int nk = num_kv_heads * embed_dim;
|
|
for (int i = threadIdx.x; i < nk; i += blockDim.x) {
|
|
const int head_idx = i / embed_dim;
|
|
const int64_t token_head =
|
|
token_idx * key_stride + head_idx * head_stride;
|
|
const int rot_offset = i % embed_dim;
|
|
apply_token_rotary_embedding<scalar_t, IS_NEOX>(
|
|
key + token_head, cos_ptr, sin_ptr, rot_offset, embed_dim);
|
|
}
|
|
}
|
|
}
|
|
|
|
template <typename scalar_t, bool IS_NEOX>
|
|
__global__ void XLLM_KERNEL_ATTR(512) rotary_embedding_kernel(
|
|
const int64_t* __restrict__ positions, // [batch_size, seq_len] or
|
|
// [num_tokens]
|
|
scalar_t* __restrict__ query, // [batch_size, seq_len, num_heads,
|
|
// head_size] or [num_tokens, num_heads,
|
|
// head_size]
|
|
scalar_t* __restrict__ key, // nullptr or
|
|
// [batch_size, seq_len, num_kv_heads,
|
|
// head_size] or [num_tokens, num_kv_heads,
|
|
// head_size]
|
|
const scalar_t* __restrict__ cos_sin_cache, // [max_position, 2,
|
|
// rot_dim // 2]
|
|
const int rot_dim,
|
|
const int64_t query_stride,
|
|
const int64_t key_stride,
|
|
const int64_t head_stride,
|
|
const int num_heads,
|
|
const int num_kv_heads,
|
|
const int head_size) {
|
|
// Each thread block is responsible for one token.
|
|
const int token_idx = blockIdx.x;
|
|
int64_t pos = positions[token_idx];
|
|
const scalar_t* cache_ptr = cos_sin_cache + pos * rot_dim;
|
|
|
|
apply_rotary_embedding<scalar_t, IS_NEOX>(query,
|
|
key,
|
|
cache_ptr,
|
|
head_size,
|
|
num_heads,
|
|
num_kv_heads,
|
|
rot_dim,
|
|
token_idx,
|
|
query_stride,
|
|
key_stride,
|
|
head_stride);
|
|
}
|
|
} // namespace
|
|
|
|
namespace xllm::kernel::cuda {
|
|
|
|
// flashinfer rope ops
|
|
// void apply_rope_pos_ids_cos_sin_cache(torch::Tensor q,
|
|
// torch::Tensor k,
|
|
// torch::Tensor cos_sin_cache,
|
|
// torch::Tensor pos_ids,
|
|
// bool interleave) {
|
|
// const int64_t head_dim = cos_sin_cache.size(-1) / 2;
|
|
// q = q.view({q.size(0), -1, head_dim});
|
|
// k = k.view({k.size(0), -1, head_dim});
|
|
|
|
// FunctionFactory::get_instance().rope_func("rope").call(
|
|
// q, k, q, k, cos_sin_cache, pos_ids, interleave);
|
|
// }
|
|
|
|
void rotary_embedding(
|
|
torch::Tensor& positions, // [batch_size, seq_len] or [num_tokens]
|
|
torch::Tensor& query, // [batch_size, seq_len, num_heads * head_size] or
|
|
// [num_tokens, num_heads * head_size] or
|
|
// [batch_size, seq_len, num_heads, head_size] or
|
|
// [num_tokens, num_heads, head_size]
|
|
std::optional<torch::Tensor> key,
|
|
// null or
|
|
// [batch_size, seq_len, num_kv_heads * head_size] or
|
|
// [num_tokens, num_kv_heads * head_size] or
|
|
// [batch_size, seq_len, num_heads, head_size] or
|
|
// [num_tokens, num_heads, head_size]
|
|
// int64_t head_size,
|
|
torch::Tensor& cos_sin_cache, // [max_position, rot_dim]
|
|
bool is_neox) {
|
|
// num_tokens = batch_size * seq_len
|
|
const int positions_ndim = positions.dim();
|
|
const int query_ndim = query.dim();
|
|
// For partial rotary models, e.g. MiniMax-M2 with head_dim=128 and
|
|
// rotary_dim=64, the cache width is the rotary dimension rather than the
|
|
// physical per-head stride. When query is already shaped as
|
|
// [*, num_heads, head_size], infer the real head_size from query itself.
|
|
int64_t head_size = (query_ndim == positions_ndim + 2)
|
|
? query.size(-1)
|
|
: cos_sin_cache.size(-1);
|
|
int64_t num_tokens = positions.numel();
|
|
|
|
// Make sure num_tokens dim is consistent across positions, query, and key
|
|
CHECK(positions_ndim == 1 || positions_ndim == 2)
|
|
<< "positions must have shape [num_tokens] or [batch_size, seq_len]";
|
|
|
|
if (positions_ndim == 1) {
|
|
CHECK(query.size(0) == positions.size(0) &&
|
|
(!key.has_value() || key->size(0) == positions.size(0)))
|
|
<< "query, key and positions must have the same number of tokens";
|
|
}
|
|
if (positions_ndim == 2) {
|
|
CHECK(query.size(0) == positions.size(0) &&
|
|
(!key.has_value() || key->size(0) == positions.size(0)) &&
|
|
query.size(1) == positions.size(1) &&
|
|
(!key.has_value() || key->size(1) == positions.size(1)))
|
|
<< "query, key and positions must have the same batch_size and seq_len";
|
|
}
|
|
|
|
// Make sure head_size is valid for query and key
|
|
// hidden_size = num_heads * head_size
|
|
int query_hidden_size = query.numel() / num_tokens;
|
|
int key_hidden_size = key.has_value() ? key->numel() / num_tokens : 0;
|
|
CHECK(query_hidden_size % head_size == 0);
|
|
CHECK(key_hidden_size % head_size == 0);
|
|
|
|
// Make sure query and key have consistent number of heads
|
|
int num_heads = query_hidden_size / head_size;
|
|
int num_kv_heads = key.has_value() ? key_hidden_size / head_size : num_heads;
|
|
CHECK(num_heads % num_kv_heads == 0);
|
|
|
|
int rot_dim = cos_sin_cache.size(1);
|
|
int seq_dim_idx = positions_ndim - 1;
|
|
int64_t query_stride = query.stride(seq_dim_idx);
|
|
int64_t key_stride = key.has_value() ? key->stride(seq_dim_idx) : 0;
|
|
// Determine head stride: for [*, heads, head_size] use stride of last dim;
|
|
// for flat [*, heads*head_size], heads blocks are contiguous of size
|
|
// head_size
|
|
int64_t head_stride =
|
|
(query_ndim == positions_ndim + 2) ? query.stride(-2) : head_size;
|
|
|
|
dim3 grid(num_tokens);
|
|
dim3 block(std::min<int64_t>(num_heads * rot_dim / 2, 512));
|
|
const at::cuda::OptionalCUDAGuard device_guard(device_of(query));
|
|
const cudaStream_t stream = at::cuda::getCurrentCUDAStream();
|
|
DISPATCH_FLOATING_TYPES(
|
|
query.scalar_type(), "apply_rope_pos_ids_cos_sin_cache", [&] {
|
|
if (is_neox) {
|
|
rotary_embedding_kernel<scalar_t, true><<<grid, block, 0, stream>>>(
|
|
positions.data_ptr<int64_t>(),
|
|
query.data_ptr<scalar_t>(),
|
|
key.has_value() ? key->data_ptr<scalar_t>() : nullptr,
|
|
cos_sin_cache.data_ptr<scalar_t>(),
|
|
rot_dim,
|
|
query_stride,
|
|
key_stride,
|
|
head_stride,
|
|
num_heads,
|
|
num_kv_heads,
|
|
head_size);
|
|
} else {
|
|
rotary_embedding_kernel<scalar_t, false><<<grid, block, 0, stream>>>(
|
|
positions.data_ptr<int64_t>(),
|
|
query.data_ptr<scalar_t>(),
|
|
key.has_value() ? key->data_ptr<scalar_t>() : nullptr,
|
|
cos_sin_cache.data_ptr<scalar_t>(),
|
|
rot_dim,
|
|
query_stride,
|
|
key_stride,
|
|
head_stride,
|
|
num_heads,
|
|
num_kv_heads,
|
|
head_size);
|
|
}
|
|
});
|
|
}
|
|
|
|
} // namespace xllm::kernel::cuda
|