161 lines
6.4 KiB
C++
161 lines
6.4 KiB
C++
#include "cache.h"
|
|
#include "ops.h"
|
|
#include "core/registration.h"
|
|
|
|
#include <torch/library.h>
|
|
|
|
std::string init_cpu_threads_env(const std::string& cpu_ids);
|
|
|
|
void int8_scaled_mm(torch::Tensor& c, const torch::Tensor& a,
|
|
const torch::Tensor& b, const torch::Tensor& a_scales,
|
|
const torch::Tensor& b_scales,
|
|
const c10::optional<torch::Tensor>& bias);
|
|
|
|
void int8_scaled_mm_azp(torch::Tensor& c, const torch::Tensor& a,
|
|
const torch::Tensor& b, const torch::Tensor& a_scales,
|
|
const torch::Tensor& b_scales,
|
|
const torch::Tensor& azp_adj,
|
|
const c10::optional<torch::Tensor>& azp,
|
|
const c10::optional<torch::Tensor>& bias);
|
|
|
|
TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) {
|
|
// vLLM custom ops
|
|
|
|
// Attention ops
|
|
// Compute the attention between an input query and the cached keys/values
|
|
// using PagedAttention.
|
|
ops.def(
|
|
"paged_attention_v1("
|
|
" Tensor! out, Tensor query, Tensor key_cache,"
|
|
" Tensor value_cache, int num_kv_heads, float scale,"
|
|
" Tensor block_tables, Tensor seq_lens, int block_size,"
|
|
" int max_seq_len, Tensor? alibi_slopes,"
|
|
" str kv_cache_dtype, float k_scale, float v_scale,"
|
|
" int tp_rank, int blocksparse_local_blocks,"
|
|
" int blocksparse_vert_stride, int blocksparse_block_size,"
|
|
" int blocksparse_head_sliding_step) -> ()");
|
|
ops.impl("paged_attention_v1", torch::kCPU, &paged_attention_v1);
|
|
|
|
// PagedAttention V2.
|
|
ops.def(
|
|
"paged_attention_v2("
|
|
" Tensor! out, Tensor! exp_sums, Tensor! max_logits,"
|
|
" Tensor! tmp_out, Tensor query, Tensor key_cache,"
|
|
" Tensor value_cache, int num_kv_heads, float scale,"
|
|
" Tensor block_tables, Tensor seq_lens, int block_size,"
|
|
" int max_seq_len, Tensor? alibi_slopes,"
|
|
" str kv_cache_dtype, float k_scale, float v_scale,"
|
|
" int tp_rank, int blocksparse_local_blocks,"
|
|
" int blocksparse_vert_stride, int blocksparse_block_size,"
|
|
" int blocksparse_head_sliding_step) -> ()");
|
|
ops.impl("paged_attention_v2", torch::kCPU, &paged_attention_v2);
|
|
|
|
// Activation ops
|
|
|
|
// Activation function used in SwiGLU.
|
|
ops.def("silu_and_mul(Tensor! out, Tensor input) -> ()");
|
|
ops.impl("silu_and_mul", torch::kCPU, &silu_and_mul);
|
|
|
|
// Activation function used in GeGLU with `none` approximation.
|
|
ops.def("gelu_and_mul(Tensor! out, Tensor input) -> ()");
|
|
ops.impl("gelu_and_mul", torch::kCPU, &gelu_and_mul);
|
|
|
|
// Activation function used in GeGLU with `tanh` approximation.
|
|
ops.def("gelu_tanh_and_mul(Tensor! out, Tensor input) -> ()");
|
|
ops.impl("gelu_tanh_and_mul", torch::kCPU, &gelu_tanh_and_mul);
|
|
|
|
// GELU implementation used in GPT-2.
|
|
ops.def("gelu_new(Tensor! out, Tensor input) -> ()");
|
|
ops.impl("gelu_new", torch::kCPU, &gelu_new);
|
|
|
|
// Approximate GELU implementation.
|
|
ops.def("gelu_fast(Tensor! out, Tensor input) -> ()");
|
|
ops.impl("gelu_fast", torch::kCPU, &gelu_fast);
|
|
|
|
// Quick GELU implementation.
|
|
ops.def("gelu_quick(Tensor! out, Tensor input) -> ()");
|
|
ops.impl("gelu_quick", torch::kCPU, &gelu_quick);
|
|
|
|
// Layernorm
|
|
// Apply Root Mean Square (RMS) Normalization to the input tensor.
|
|
ops.def(
|
|
"rms_norm(Tensor! out, Tensor input, Tensor weight, float epsilon) -> "
|
|
"()");
|
|
ops.impl("rms_norm", torch::kCPU, &rms_norm);
|
|
|
|
// In-place fused Add and RMS Normalization.
|
|
ops.def(
|
|
"fused_add_rms_norm(Tensor! input, Tensor! residual, Tensor weight, "
|
|
"float epsilon) -> ()");
|
|
ops.impl("fused_add_rms_norm", torch::kCPU, &fused_add_rms_norm);
|
|
|
|
// Rotary embedding
|
|
// Apply GPT-NeoX or GPT-J style rotary embedding to query and key.
|
|
ops.def(
|
|
"rotary_embedding(Tensor positions, Tensor! query,"
|
|
" Tensor! key, int head_size,"
|
|
" Tensor cos_sin_cache, bool is_neox) -> ()");
|
|
ops.impl("rotary_embedding", torch::kCPU, &rotary_embedding);
|
|
|
|
// Quantization
|
|
#ifdef __AVX512F__
|
|
// Compute int8 quantized tensor for given scaling factor.
|
|
ops.def(
|
|
"static_scaled_int8_quant(Tensor! out, Tensor input, Tensor scale,"
|
|
"Tensor? azp) -> ()");
|
|
ops.impl("static_scaled_int8_quant", torch::kCPU, &static_scaled_int8_quant);
|
|
|
|
// Compute int8 quantized tensor and scaling factor
|
|
ops.def(
|
|
"dynamic_scaled_int8_quant(Tensor! out, Tensor input, Tensor! scale, "
|
|
"Tensor!? azp) -> ()");
|
|
ops.impl("dynamic_scaled_int8_quant", torch::kCPU,
|
|
&dynamic_scaled_int8_quant);
|
|
// W8A8 GEMM, supporting symmetric per-tensor or per-row/column
|
|
// quantization.
|
|
ops.def(
|
|
"cutlass_scaled_mm(Tensor! out, Tensor a,"
|
|
" Tensor b, Tensor a_scales,"
|
|
" Tensor b_scales, Tensor? bias) -> ()");
|
|
ops.impl("cutlass_scaled_mm", torch::kCPU, &int8_scaled_mm);
|
|
// w8a8 GEMM, supporting asymmetric per-tensor or per-row/column
|
|
// quantization.
|
|
ops.def(
|
|
"cutlass_scaled_mm_azp(Tensor! out, Tensor a,"
|
|
" Tensor b, Tensor a_scales,"
|
|
" Tensor b_scales, Tensor azp_adj,"
|
|
" Tensor? azp, Tensor? bias) -> ()");
|
|
ops.impl("cutlass_scaled_mm_azp", torch::kCPU, &int8_scaled_mm_azp);
|
|
#endif
|
|
}
|
|
|
|
TORCH_LIBRARY_EXPAND(CONCAT(TORCH_EXTENSION_NAME, _cache_ops), cache_ops) {
|
|
// Cache ops
|
|
// Swap in (out) the cache blocks from src to dst.
|
|
cache_ops.def(
|
|
"swap_blocks(Tensor src, Tensor! dst, Tensor block_mapping) -> ()");
|
|
cache_ops.impl("swap_blocks", torch::kCPU, &swap_blocks);
|
|
|
|
// Copy the cache blocks from src to dst.
|
|
cache_ops.def(
|
|
"copy_blocks(Tensor(a!)[] key_caches, Tensor[](b!) value_caches, "
|
|
"Tensor block_mapping) -> ()");
|
|
cache_ops.impl("copy_blocks", torch::kCPU, ©_blocks);
|
|
|
|
// Reshape the key and value tensors and cache them.
|
|
cache_ops.def(
|
|
"reshape_and_cache(Tensor key, Tensor value,"
|
|
" Tensor! key_cache, Tensor! value_cache,"
|
|
" Tensor slot_mapping,"
|
|
" str kv_cache_dtype,"
|
|
" float k_scale, float v_scale) -> ()");
|
|
cache_ops.impl("reshape_and_cache", torch::kCPU, &reshape_and_cache);
|
|
}
|
|
|
|
TORCH_LIBRARY_EXPAND(CONCAT(TORCH_EXTENSION_NAME, _utils), utils) {
|
|
// CPU utils
|
|
utils.def("init_cpu_threads_env(str cpu_ids) -> str", &init_cpu_threads_env);
|
|
}
|
|
|
|
REGISTER_EXTENSION(TORCH_EXTENSION_NAME)
|