@tduncan / Nvidia-Optimization-Stack-VLLM
Code Commits Issues Pulls Cloud Agent CI/CD
main
main
h 72 lines 2.86 KB
Raw
1 #pragma once
2
3 #include <optional>
4 #include <string>
5 #include <torch/library.h>
6 #include <tuple>
7
8 #include "core/scalar_type.hpp"
9
10 #include <vector>
11
12 // rms_norm and fused_add_rms_norm declarations also exist in
13 // csrc/libtorch_stable/ops.h (torch::stable ABI for CUDA). They remain here
14 // because the CPU build still uses these torch::Tensor declarations.
15 void rms_norm(torch::Tensor& out, torch::Tensor& input,
16 std::optional<torch::Tensor> weight, double epsilon);
17
18 void fused_add_rms_norm(torch::Tensor& input, torch::Tensor& residual,
19 std::optional<torch::Tensor> weight, double epsilon);
20
21 // rotary_embedding also exist in csrc/libtorch_stable/ops.h (torch::stable
22 // ABI for CUDA). It remains here because the CPU build still uses these
23 // torch::Tensor declarations.
24 void rotary_embedding(torch::Tensor& positions, torch::Tensor& query,
25 std::optional<torch::Tensor> key, int64_t head_size,
26 torch::Tensor& cos_sin_cache, bool is_neox,
27 int64_t rope_dim_offset, bool inverse);
28
29 void silu_and_mul(torch::Tensor& out, torch::Tensor& input);
30
31 void silu_and_mul_clamp(torch::Tensor& out, torch::Tensor& input, double limit,
32 double alpha = 1.0, double beta = 0.0);
33
34 void gelu_and_mul(torch::Tensor& out, torch::Tensor& input);
35
36 void gelu_tanh_and_mul(torch::Tensor& out, torch::Tensor& input);
37
38 void gelu_tanh(torch::Tensor& out, torch::Tensor& input);
39
40 void gelu_new(torch::Tensor& out, torch::Tensor& input);
41
42 void gelu_fast(torch::Tensor& out, torch::Tensor& input);
43
44 void gelu_quick(torch::Tensor& out, torch::Tensor& input);
45
46 void relu_squared(torch::Tensor& out, torch::Tensor& input);
47
48 void static_scaled_int8_quant(torch::Tensor& out, torch::Tensor const& input,
49 torch::Tensor const& scale,
50 std::optional<torch::Tensor> const& azp);
51
52 void dynamic_scaled_int8_quant(torch::Tensor& out, torch::Tensor const& input,
53 torch::Tensor& scales,
54 std::optional<torch::Tensor> const& azp);
55
56 torch::Tensor dynamic_4bit_int_moe_cpu(
57 torch::Tensor x, torch::Tensor topk_ids, torch::Tensor topk_weights,
58 torch::Tensor w13_packed, torch::Tensor w2_packed, int64_t hidden_size,
59 int64_t intermediate_size, int64_t group_size,
60 bool apply_router_weight_on_input, int64_t activation_kind);
61
62 using fptr_t = int64_t;
63 #ifdef USE_ROCM
64 fptr_t init_custom_qr(int64_t rank, int64_t world_size,
65 std::optional<int64_t> qr_max_size = std::nullopt);
66 void qr_destroy(fptr_t _fa);
67 torch::Tensor qr_get_handle(fptr_t _fa);
68 void qr_open_handles(fptr_t _fa, const std::vector<torch::Tensor>& handles);
69 void qr_all_reduce(fptr_t _fa, torch::Tensor& inp, torch::Tensor& out,
70 int64_t quant_level, bool cast_bf2half = false);
71 int64_t qr_max_size();
72 #endif