@tduncan / Nvidia-Optimization-Stack-VLLM
Code Commits Issues Pulls Cloud Agent CI/CD
main
main
cpp 136 lines 4.71 KB
Raw
1 #include "cpu_types.hpp"
2
3 namespace {
4 template <typename scalar_t>
5 void rms_norm_impl(scalar_t* __restrict__ out,
6 const scalar_t* __restrict__ input,
7 const scalar_t* __restrict__ weight, const bool has_weight,
8 const float epsilon, const int num_tokens,
9 const int hidden_size) {
10 using scalar_vec_t = vec_op::vec_t<scalar_t>;
11 constexpr int VEC_ELEM_NUM = scalar_vec_t::get_elem_num();
12 TORCH_CHECK(hidden_size % VEC_ELEM_NUM == 0);
13
14 #pragma omp parallel for
15 for (int i = 0; i < num_tokens; ++i) {
16 vec_op::FP32Vec8 variance(0.0);
17 auto input_p = input + i * hidden_size;
18 auto output_p = out + i * hidden_size;
19 for (int j = 0; j < hidden_size; j += VEC_ELEM_NUM) {
20 scalar_vec_t x(input_p + j);
21 vec_op::FP32Vec8 fp32_x(x);
22 variance = variance + fp32_x * fp32_x;
23 }
24
25 float s_variance =
26 1.0f / sqrtf(variance.reduce_sum() / (float)hidden_size + epsilon);
27 vec_op::FP32Vec8 fp32_s_variance(s_variance);
28
29 for (int j = 0; j < hidden_size; j += VEC_ELEM_NUM) {
30 scalar_vec_t x(input_p + j);
31 vec_op::FP32Vec8 fp32_x(x);
32 vec_op::FP32Vec8 fp32_out;
33 if (has_weight) {
34 scalar_vec_t w(weight + j);
35 vec_op::FP32Vec8 fp32_w(w);
36 fp32_out = fp32_x * fp32_s_variance * fp32_w;
37 } else {
38 fp32_out = fp32_x * fp32_s_variance;
39 }
40
41 scalar_vec_t out(fp32_out);
42 out.save(output_p + j);
43 }
44 }
45 }
46
47 template <typename scalar_t>
48 void fused_add_rms_norm_impl(scalar_t* __restrict__ input,
49 scalar_t* __restrict__ residual,
50 const scalar_t* __restrict__ weight,
51 const bool has_weight, const float epsilon,
52 const int num_tokens, const int hidden_size) {
53 using scalar_vec_t = vec_op::vec_t<scalar_t>;
54 constexpr int VEC_ELEM_NUM = scalar_vec_t::get_elem_num();
55 TORCH_CHECK(hidden_size % VEC_ELEM_NUM == 0);
56
57 #pragma omp parallel for
58 for (int i = 0; i < num_tokens; ++i) {
59 vec_op::FP32Vec8 variance(0.0);
60 auto input_p = input + i * hidden_size;
61 auto residual_p = residual + i * hidden_size;
62 for (int j = 0; j < hidden_size; j += VEC_ELEM_NUM) {
63 scalar_vec_t x(input_p + j);
64 scalar_vec_t res(residual_p + j);
65 vec_op::FP32Vec8 fp32_x(x);
66 vec_op::FP32Vec8 fp32_res(res);
67
68 fp32_x = fp32_x + fp32_res;
69 variance = variance + fp32_x * fp32_x;
70 scalar_vec_t out(fp32_x);
71 out.save(residual_p + j);
72 }
73
74 float s_variance =
75 1.0f / sqrtf(variance.reduce_sum() / (float)hidden_size + epsilon);
76 vec_op::FP32Vec8 fp32_s_variance(s_variance);
77
78 for (int j = 0; j < hidden_size; j += VEC_ELEM_NUM) {
79 vec_op::FP32Vec8 fp32_out;
80 if (has_weight) {
81 scalar_vec_t w(weight + j);
82 scalar_vec_t res(residual_p + j);
83 vec_op::FP32Vec8 fp32_w(w);
84 vec_op::FP32Vec8 fp32_res(res);
85 fp32_out = fp32_res * fp32_s_variance * fp32_w;
86 } else {
87 scalar_vec_t res(residual_p + j);
88 vec_op::FP32Vec8 fp32_res(res);
89 fp32_out = fp32_res * fp32_s_variance;
90 }
91
92 scalar_vec_t out(fp32_out);
93 out.save(input_p + j);
94 }
95 }
96 }
97 } // namespace
98
99 void rms_norm(torch::Tensor& out, torch::Tensor& input,
100 std::optional<torch::Tensor> weight, double epsilon) {
101 int hidden_size = input.size(-1);
102 int num_tokens = input.numel() / hidden_size;
103 const bool has_weight = weight.has_value();
104 if (has_weight) {
105 TORCH_CHECK(weight->is_contiguous());
106 }
107
108 VLLM_DISPATCH_FLOATING_TYPES(input.scalar_type(), "rms_norm_impl", [&] {
109 CPU_KERNEL_GUARD_IN(rms_norm_impl)
110 rms_norm_impl(out.data_ptr<scalar_t>(), input.data_ptr<scalar_t>(),
111 has_weight ? weight->data_ptr<scalar_t>() : nullptr,
112 has_weight, epsilon, num_tokens, hidden_size);
113 CPU_KERNEL_GUARD_OUT(rms_norm_impl)
114 });
115 }
116
117 void fused_add_rms_norm(torch::Tensor& input, torch::Tensor& residual,
118 std::optional<torch::Tensor> weight, double epsilon) {
119 int hidden_size = input.size(-1);
120 int num_tokens = input.numel() / hidden_size;
121 const bool has_weight = weight.has_value();
122 if (has_weight) {
123 TORCH_CHECK(weight->scalar_type() == input.scalar_type());
124 TORCH_CHECK(weight->is_contiguous());
125 }
126
127 VLLM_DISPATCH_FLOATING_TYPES(
128 input.scalar_type(), "fused_add_rms_norm_impl", [&] {
129 CPU_KERNEL_GUARD_IN(fused_add_rms_norm_impl)
130 fused_add_rms_norm_impl(
131 input.data_ptr<scalar_t>(), residual.data_ptr<scalar_t>(),
132 has_weight ? weight->data_ptr<scalar_t>() : nullptr, has_weight,
133 epsilon, num_tokens, hidden_size);
134 CPU_KERNEL_GUARD_OUT(fused_add_rms_norm_impl)
135 });
136 }