A high-throughput and memory-efficient inference and serving engine for LLMs
Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git
Synced Nov 14, 05:30 UTC
[Chore] Remove unused `PolyNorm` layer (#27110)
Signed-off-by: Isotr0py <mozf@mail2.sysu.edu.cn>
Isotr0py committed
Oct 18, 2025 at 03:03 UTC
3125d7995020f4ae1b7777e41af6c3a6e714cb1c
7 files changed
+1
-524
benchmarks/kernels/benchmark_polynorm.py
deleted
-155
@@ -1,155 +0,0 @@
1
-# SPDX-License-Identifier: Apache-2.0
2
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
-
4
-import itertools
5
-
6
-import torch
7
-
8
-from vllm import _custom_ops as vllm_ops
9
-from vllm.triton_utils import triton
10
-
11
-
12
-def polynorm_naive(
13
- x: torch.Tensor,
14
- weight: torch.Tensor,
15
- bias: torch.Tensor,
16
- eps: float = 1e-6,
17
-):
18
- orig_shape = x.shape
19
- x = x.view(-1, x.shape[-1])
20
-
21
- def norm(x, eps: float):
22
- return x / torch.sqrt(x.pow(2).mean(-1, keepdim=True) + eps)
23
-
24
- x = x.float()
25
- return (
26
- (
27
- weight[0] * norm(x**3, eps)
28
- + weight[1] * norm(x**2, eps)
29
- + weight[2] * norm(x, eps)
30
- + bias
31
- )
32
- .to(weight.dtype)
33
- .view(orig_shape)
34
- )
35
-
36
-
37
-def polynorm_vllm(
38
- x: torch.Tensor,
39
- weight: torch.Tensor,
40
- bias: torch.Tensor,
41
- eps: float = 1e-6,
42
-):
43
- orig_shape = x.shape
44
- x = x.view(-1, x.shape[-1])
45
-
46
- out = torch.empty_like(x)
47
- vllm_ops.poly_norm(out, x, weight, bias, eps)
48
- output = out
49
-
50
- output = output.view(orig_shape)
51
- return output
52
-
53
-
54
-def calculate_diff(batch_size, seq_len, hidden_dim):
55
- dtype = torch.bfloat16
56
- x = torch.randn(batch_size, seq_len, hidden_dim, dtype=dtype, device="cuda")
57
- weight = torch.ones(3, dtype=dtype, device="cuda")
58
- bias = torch.ones(1, dtype=dtype, device="cuda")
59
-
60
- output_naive = polynorm_naive(x, weight, bias)
61
- output_vllm = polynorm_vllm(x, weight, bias)
62
-
63
- if torch.allclose(output_naive, output_vllm, atol=1e-2, rtol=1e-2):
64
- print("✅ All implementations match")
65
- else:
66
- print("❌ Implementations differ")
67
-
68
-
69
-batch_size_range = [2**i for i in range(0, 7, 2)]
70
-seq_length_range = [2**i for i in range(6, 11, 1)]
71
-dim_range = [2048, 4096]
72
-configs = list(itertools.product(dim_range, batch_size_range, seq_length_range))
73
-
74
-
75
-def get_benchmark():
76
- @triton.testing.perf_report(
77
- triton.testing.Benchmark(
78
- x_names=["dim", "batch_size", "seq_len"],
79
- x_vals=[list(_) for _ in configs],
80
- line_arg="provider",
81
- line_vals=["naive", "vllm"],
82
- line_names=["Naive", "vLLM"],
83
- styles=[("blue", "-"), ("red", "-")],
84
- ylabel="us",
85
- plot_name="polynorm-perf",
86
- args={},
87
- )
88
- )
89
- def benchmark(dim, batch_size, seq_len, provider):
90
- dtype = torch.bfloat16
91
- hidden_dim = dim * 4
92
-
93
- x = torch.randn(batch_size, seq_len, hidden_dim, dtype=dtype, device="cuda")
94
- weight = torch.ones(3, dtype=dtype, device="cuda")
95
- bias = torch.ones(1, dtype=dtype, device="cuda")
96
-
97
- quantiles = [0.5, 0.2, 0.8]
98
-
99
- if provider == "naive":
100
- ms, min_ms, max_ms = triton.testing.do_bench(
101
- lambda: polynorm_naive(x, weight, bias),
102
- quantiles=quantiles,
103
- )
104
- else:
105
- ms, min_ms, max_ms = triton.testing.do_bench(
106
- lambda: polynorm_vllm(x, weight, bias),
107
- quantiles=quantiles,
108
- )
109
-
110
- return 1000 * ms, 1000 * max_ms, 1000 * min_ms
111
-
112
- return benchmark
113
-
114
-
115
-if __name__ == "__main__":
116
- import argparse
117
-
118
- parser = argparse.ArgumentParser()
119
- parser.add_argument(
120
- "--batch-size",
121
- type=int,
122
- default=4,
123
- help="Batch size",
124
- )
125
- parser.add_argument(
126
- "--seq-len",
127
- type=int,
128
- default=128,
129
- help="Sequence length",
130
- )
131
- parser.add_argument(
132
- "--hidden-dim",
133
- type=int,
134
- default=8192,
135
- help="Intermediate size of MLP",
136
- )
137
- parser.add_argument(
138
- "--save-path",
139
- type=str,
140
- default="./configs/polnorm/",
141
- help="Path to save polnorm benchmark results",
142
- )
143
-
144
- args = parser.parse_args()
145
-
146
- # Run correctness test
147
- calculate_diff(
148
- batch_size=args.batch_size,
149
- seq_len=args.seq_len,
150
- hidden_dim=args.hidden_dim,
151
- )
152
-
153
- benchmark = get_benchmark()
154
- # Run performance benchmark
155
- benchmark.run(print_data=True, save_path=args.save_path)
csrc/layernorm_kernels.cu
-252
@@ -148,211 +148,6 @@ fused_add_rms_norm_kernel(
148
}
149
}
150
151
-/* Function specialization in the case of FP16/BF16 tensors.
152
- Additional optimizations we can make in this case are
153
- packed and vectorized operations, which help with the
154
- memory latency bottleneck.
155
-
156
- _f16VecPN struct extends _f16Vec to add operations specifically required for
157
- polynomial normalization (poly norm).
158
- The original _f16Vec does not include the sum-of-powers computation or
159
- in-place polynomial normalization logic. */
160
-template <typename scalar_t, int width>
161
-struct alignas(16) _f16VecPN : _f16Vec<scalar_t, width> {
162
- using Base = _f16Vec<scalar_t, width>;
163
- using Converter = typename Base::Converter;
164
- using T1 = typename Base::T1;
165
- using T2 = typename Base::T2;
166
- using Base::data;
167
-
168
- __device__ auto sum_pows() const {
169
- float s2 = 0.0f, s4 = 0.0f, s6 = 0.0f;
170
-
171
-#pragma unroll
172
- for (int i = 0; i < width; i += 2) {
173
- float2 z = Converter::convert(T2{data[i], data[i + 1]});
174
- float x2 = z.x * z.x;
175
- float x4 = x2 * x2;
176
- float x6 = x4 * x2;
177
-
178
- float y2 = z.y * z.y;
179
- float y4 = y2 * y2;
180
- float y6 = y4 * y2;
181
-
182
- s2 += x2 + y2;
183
- s4 += x4 + y4;
184
- s6 += x6 + y6;
185
- }
186
- return std::make_tuple(s2, s4, s6);
187
- }
188
-
189
- __device__ void poly_norm_inplace(const float w2_inv_std,
190
- const float w1_inv_std2,
191
- const float w0_inv_std3, const float bias) {
192
-#pragma unroll
193
- for (int i = 0; i < width; i += 2) {
194
- float2 z = Converter::convert(T2{data[i], data[i + 1]});
195
-
196
- float x2 = z.x * z.x;
197
- float x3 = x2 * z.x;
198
- z.x = w2_inv_std * z.x + w1_inv_std2 * x2 + w0_inv_std3 * x3 + bias;
199
-
200
- float y2 = z.y * z.y;
201
- float y3 = y2 * z.y;
202
- z.y = w2_inv_std * z.y + w1_inv_std2 * y2 + w0_inv_std3 * y3 + bias;
203
-
204
- auto out = Converter::convert(z);
205
- data[i] = out.x;
206
- data[i + 1] = out.y;
207
- }
208
- }
209
-};
210
-
211
-template <typename scalar_t, int width>
212
-__global__ std::enable_if_t<(width > 0) && _typeConvert<scalar_t>::exists>
213
-poly_norm_kernel(scalar_t* __restrict__ out, // [..., hidden_size]
214
- const scalar_t* __restrict__ input, // [..., hidden_size]
215
- const scalar_t* __restrict__ weight, // [3]
216
- const scalar_t* __restrict__ bias, // [1]
217
- const float epsilon, const int hidden_size) {
218
- // Sanity checks on our vector struct and type-punned pointer arithmetic
219
- static_assert(std::is_pod_v<_f16VecPN<scalar_t, width>>);
220
- static_assert(sizeof(_f16VecPN<scalar_t, width>) == sizeof(scalar_t) * width);
221
-
222
- /* These and the argument pointers are all declared `restrict` as they are
223
- not aliased in practice. Argument pointers should not be dereferenced
224
- in this kernel as that would be undefined behavior */
225
- auto* __restrict__ input_v =
226
- reinterpret_cast<const _f16VecPN<scalar_t, width>*>(input);
227
- const int vec_hidden_size = hidden_size / width;
228
- float variance = 0.0f;
229
- float variance2 = 0.0f;
230
- float variance3 = 0.0f;
231
-
232
- for (int idx = threadIdx.x; idx < vec_hidden_size; idx += blockDim.x) {
233
- int id = blockIdx.x * vec_hidden_size + idx;
234
- _f16VecPN<scalar_t, width> temp = input_v[id];
235
- auto [x2, x4, x6] = temp.sum_pows();
236
-
237
- variance += x2;
238
- variance2 += x4;
239
- variance3 += x6;
240
- }
241
-
242
- float3 thread_variances = make_float3(variance, variance2, variance3);
243
-
244
- struct SumOp {
245
- __device__ float3 operator()(const float3& a, const float3& b) const {
246
- return make_float3(a.x + b.x, a.y + b.y, a.z + b.z);
247
- }
248
- };
249
-
250
- using BlockReduce = cub::BlockReduce<float3, 1024>;
251
- __shared__ typename BlockReduce::TempStorage reduceStore;
252
- float3 block_variances =
253
- BlockReduce(reduceStore).Reduce(thread_variances, SumOp{}, blockDim.x);
254
-
255
- variance = block_variances.x;
256
- variance2 = block_variances.y;
257
- variance3 = block_variances.z;
258
-
259
- __shared__ float s_w2_inv_std;
260
- __shared__ float s_w1_inv_std2;
261
- __shared__ float s_w0_inv_std3;
262
- __shared__ float s_bias;
263
-
264
- if (threadIdx.x == 0) {
265
- float w0 = (float)weight[0];
266
- float w1 = (float)weight[1];
267
- float w2 = (float)weight[2];
268
- s_bias = (float)bias[0];
269
-
270
- s_w2_inv_std = w2 * rsqrtf(variance / hidden_size + epsilon);
271
- s_w1_inv_std2 = w1 * rsqrtf(variance2 / hidden_size + epsilon);
272
- s_w0_inv_std3 = w0 * rsqrtf(variance3 / hidden_size + epsilon);
273
- }
274
- __syncthreads();
275
-
276
- auto* __restrict__ out_v = reinterpret_cast<_f16VecPN<scalar_t, width>*>(out);
277
-
278
- for (int idx = threadIdx.x; idx < vec_hidden_size; idx += blockDim.x) {
279
- int id = blockIdx.x * vec_hidden_size + idx;
280
- _f16VecPN<scalar_t, width> temp = input_v[id];
281
- temp.poly_norm_inplace(s_w2_inv_std, s_w1_inv_std2, s_w0_inv_std3, s_bias);
282
- out_v[id] = temp;
283
- }
284
-}
285
-
286
-/* Generic poly_norm_kernel
287
- The width field is not used here but necessary for other specializations.
288
- */
289
-template <typename scalar_t, int width>
290
-__global__ std::enable_if_t<(width == 0) || !_typeConvert<scalar_t>::exists>
291
-poly_norm_kernel(scalar_t* __restrict__ out, // [..., hidden_size]
292
- const scalar_t* __restrict__ input, // [..., hidden_size]
293
- const scalar_t* __restrict__ weight, // [3]
294
- const scalar_t* __restrict__ bias, // [1]
295
- const float epsilon, const int hidden_size) {
296
- float variance = 0.0f;
297
- float variance2 = 0.0f;
298
- float variance3 = 0.0f;
299
-
300
- for (int idx = threadIdx.x; idx < hidden_size; idx += blockDim.x) {
301
- float x = (float)input[blockIdx.x * hidden_size + idx];
302
- float x2 = x * x;
303
- float x4 = x2 * x2;
304
- float x6 = x4 * x2;
305
-
306
- variance += x2;
307
- variance2 += x4;
308
- variance3 += x6;
309
- }
310
-
311
- float3 thread_variances = make_float3(variance, variance2, variance3);
312
-
313
- struct SumOp {
314
- __device__ float3 operator()(const float3& a, const float3& b) const {
315
- return make_float3(a.x + b.x, a.y + b.y, a.z + b.z);
316
- }
317
- };
318
-
319
- using BlockReduce = cub::BlockReduce<float3, 1024>;
320
- __shared__ typename BlockReduce::TempStorage reduceStore;
321
- float3 block_variances =
322
- BlockReduce(reduceStore).Reduce(thread_variances, SumOp{}, blockDim.x);
323
-
324
- variance = block_variances.x;
325
- variance2 = block_variances.y;
326
- variance3 = block_variances.z;
327
-
328
- __shared__ float s_w2_inv_std;
329
- __shared__ float s_w1_inv_std2;
330
- __shared__ float s_w0_inv_std3;
331
- __shared__ float s_bias;
332
-
333
- if (threadIdx.x == 0) {
334
- float w0 = (float)weight[0];
335
- float w1 = (float)weight[1];
336
- float w2 = (float)weight[2];
337
- s_bias = (float)bias[0];
338
-
339
- s_w2_inv_std = w2 * rsqrtf(variance / hidden_size + epsilon);
340
- s_w1_inv_std2 = w1 * rsqrtf(variance2 / hidden_size + epsilon);
341
- s_w0_inv_std3 = w0 * rsqrtf(variance3 / hidden_size + epsilon);
342
- }
343
- __syncthreads();
344
-
345
- for (int idx = threadIdx.x; idx < hidden_size; idx += blockDim.x) {
346
- float x = (float)input[blockIdx.x * hidden_size + idx];
347
- float x2 = x * x;
348
- float x3 = x2 * x;
349
-
350
- out[blockIdx.x * hidden_size + idx] =
351
- (scalar_t)(x * s_w2_inv_std + x2 * s_w1_inv_std2 + x3 * s_w0_inv_std3 +
352
- s_bias);
353
- }
354
-}
355
-
151
} // namespace vllm
152
153
void rms_norm(torch::Tensor& out, // [..., hidden_size]
@@ -444,50 +239,3 @@ void fused_add_rms_norm(torch::Tensor& input, // [..., hidden_size]
239
LAUNCH_FUSED_ADD_RMS_NORM(0);
240
}
241
}
447
-
448
-#define LAUNCH_FUSED_POLY_NORM(width) \
449
- VLLM_DISPATCH_FLOATING_TYPES(input.scalar_type(), "poly_norm_kernel", [&] { \
450
- vllm::poly_norm_kernel<scalar_t, width><<<grid, block, 0, stream>>>( \
451
- out.data_ptr<scalar_t>(), input.data_ptr<scalar_t>(), \
452
- weight.data_ptr<scalar_t>(), bias.data_ptr<scalar_t>(), epsilon, \
453
- hidden_size); \
454
- });
455
-
456
-void poly_norm(torch::Tensor& out, // [..., hidden_size]
457
- torch::Tensor& input, // [..., hidden_size]
458
- torch::Tensor& weight, // [3]
459
- torch::Tensor& bias, // [1]
460
- double epsilon) {
461
- TORCH_CHECK(out.is_contiguous());
462
- TORCH_CHECK(input.is_contiguous());
463
- TORCH_CHECK(out.data_ptr() != input.data_ptr());
464
-
465
- int hidden_size = input.size(-1);
466
- int num_tokens = input.numel() / hidden_size;
467
-
468
- dim3 grid(num_tokens);
469
- /* This kernel is memory-latency bound in many scenarios.
470
- When num_tokens is large, a smaller block size allows
471
- for increased block occupancy on CUs and better latency
472
- hiding on global mem ops. */
473
- const int max_block_size = (num_tokens < 256) ? 1024 : 256;
474
- dim3 block(std::min(hidden_size, max_block_size));
475
- const at::cuda::OptionalCUDAGuard device_guard(device_of(input));
476
- const cudaStream_t stream = at::cuda::getCurrentCUDAStream();
477
- /*If the tensor types are FP16/BF16, try to use the optimized kernel
478
- with packed + vectorized ops.
479
- Max optimization is achieved with a width-8 vector of FP16/BF16s
480
- since we can load at most 128 bits at once in a global memory op.
481
- However, this requires each tensor's data to be aligned to 16
482
- bytes.
483
- */
484
- auto inp_ptr = reinterpret_cast<std::uintptr_t>(input.data_ptr());
485
- auto out_ptr = reinterpret_cast<std::uintptr_t>(out.data_ptr());
486
- bool ptrs_are_aligned = inp_ptr % 16 == 0 && out_ptr % 16 == 0;
487
- bool batch_invariant_launch = vllm::vllm_is_batch_invariant();
488
- if (ptrs_are_aligned && hidden_size % 8 == 0 && !batch_invariant_launch) {
489
- LAUNCH_FUSED_POLY_NORM(8);
490
- } else {
491
- LAUNCH_FUSED_POLY_NORM(0);
492
- }
493
-}
csrc/ops.h
-3
@@ -92,9 +92,6 @@ void rms_norm(torch::Tensor& out, torch::Tensor& input, torch::Tensor& weight,
92
void fused_add_rms_norm(torch::Tensor& input, torch::Tensor& residual,
93
torch::Tensor& weight, double epsilon);
94
95
-void poly_norm(torch::Tensor& out, torch::Tensor& input, torch::Tensor& weight,
96
- torch::Tensor& bias, double epsilon);
97
-
95
void apply_repetition_penalties_(torch::Tensor& logits,
96
const torch::Tensor& prompt_mask,
97
const torch::Tensor& output_mask,
csrc/torch_bindings.cpp
-6
@@ -175,12 +175,6 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) {
175
"float epsilon) -> ()");
176
ops.impl("fused_add_rms_norm", torch::kCUDA, &fused_add_rms_norm);
177
178
- // Polynomial Normalization.
179
- ops.def(
180
- "poly_norm(Tensor! out, Tensor input, Tensor weight, Tensor bias, float "
181
- "epsilon) -> ()");
182
- ops.impl("poly_norm", torch::kCUDA, &poly_norm);
183
-
178
// Apply repetition penalties to logits in-place
179
ops.def(
180
"apply_repetition_penalties_(Tensor! logits, Tensor prompt_mask, "
tests/kernels/core/test_layernorm.py
+1
-33
@@ -6,7 +6,7 @@ import torch
6
7
from tests.kernels.quant_utils import FP8_DTYPE
8
from tests.kernels.utils import opcheck
9
-from vllm.model_executor.layers.layernorm import PolyNorm, RMSNorm
9
+from vllm.model_executor.layers.layernorm import RMSNorm
10
from vllm.platforms import current_platform
11
12
DTYPES = [torch.half, torch.bfloat16, torch.float]
@@ -70,38 +70,6 @@ def test_rms_norm(
70
)
71
72
73
-@pytest.mark.parametrize("num_tokens", NUM_TOKENS)
74
-@pytest.mark.parametrize("hidden_size", HIDDEN_SIZES)
75
-@pytest.mark.parametrize("dtype", DTYPES)
76
-@pytest.mark.parametrize("seed", SEEDS)
77
-@pytest.mark.parametrize("device", CUDA_DEVICES)
78
-@torch.inference_mode()
79
-def test_poly_norm(
80
- num_tokens: int,
81
- hidden_size: int,
82
- dtype: torch.dtype,
83
- seed: int,
84
- device: str,
85
-) -> None:
86
- current_platform.seed_everything(seed)
87
- torch.set_default_device(device)
88
- layer = PolyNorm().to(dtype=dtype)
89
- layer.weight.data.normal_(mean=1.0, std=0.1)
90
- layer.bias.data.normal_(mean=1.0, std=0.1)
91
- scale = 1 / (2 * hidden_size)
92
- x = torch.randn(num_tokens, hidden_size, dtype=dtype)
93
- x *= scale
94
-
95
- ref_out = layer.forward_native(x)
96
- out = layer(x)
97
- torch.testing.assert_close(out, ref_out, atol=1e-2, rtol=1e-2)
98
-
99
- opcheck(
100
- torch.ops._C.poly_norm,
101
- (out, x, layer.weight.data, layer.bias.data, layer.variance_epsilon),
102
- )
103
-
104
-
73
@pytest.mark.parametrize("num_tokens", NUM_TOKENS)
74
@pytest.mark.parametrize("hidden_size", HIDDEN_SIZES)
75
@pytest.mark.parametrize("add_residual", ADD_RESIDUAL)
vllm/_custom_ops.py
-12
@@ -339,18 +339,6 @@ def fused_add_rms_norm(
339
torch.ops._C.fused_add_rms_norm(input, residual, weight, epsilon)
340
341
342
-def poly_norm(
343
- out: torch.Tensor,
344
- input: torch.Tensor,
345
- weight: torch.Tensor,
346
- bias: torch.Tensor,
347
- epsilon: float,
348
-) -> None:
349
- # TODO: Remove this contiguous call when the kernel is updated to support non-contiguous input
350
- input_contiguous = input.contiguous()
351
- torch.ops._C.poly_norm(out, input_contiguous, weight, bias, epsilon)
352
-
353
-
342
def apply_repetition_penalties_torch(
343
logits: torch.Tensor,
344
prompt_mask: torch.Tensor,
vllm/model_executor/layers/layernorm.py
-63
@@ -58,22 +58,6 @@ def fused_add_rms_norm(
58
return x, residual
59
60
61
-def poly_norm(
62
- x: torch.Tensor, weight: torch.Tensor, bias: torch.Tensor, variance_epsilon: float
63
-) -> torch.Tensor:
64
- from vllm import _custom_ops as ops
65
-
66
- out = torch.empty_like(x)
67
- ops.poly_norm(
68
- out,
69
- x,
70
- weight,
71
- bias,
72
- variance_epsilon,
73
- )
74
- return out
75
-
76
-
61
def rocm_aiter_rms_norm_impl(
62
x: torch.Tensor, weight: torch.Tensor, variance_epsilon: float
63
) -> torch.Tensor:
@@ -385,53 +369,6 @@ class GemmaRMSNorm(CustomOp):
369
return self.forward_native(x, residual)
370
371
388
-@CustomOp.register("poly_norm")
389
-class PolyNorm(CustomOp):
390
- """Polynomial normalization.
391
-
392
- Computes x -> w_0 * RMSNorm(x^3) + w_1 * RMSNorm(x^2) + w_2 * RMSNorm(x) + b
393
- where w_n is the learned weight and b is the bias.
394
- Refer to https://arxiv.org/html/2411.03884v1
395
- """
396
-
397
- def __init__(
398
- self,
399
- eps: float = 1e-6,
400
- ) -> None:
401
- super().__init__()
402
- self.weight = torch.nn.Parameter(torch.ones(3) / 3)
403
- self.bias = torch.nn.Parameter(torch.zeros(1))
404
- self.variance_epsilon = eps
405
-
406
- def _norm(self, x):
407
- return x / torch.sqrt(x.pow(2).mean(-1, keepdim=True) + self.variance_epsilon)
408
-
409
- def forward_native(
410
- self,
411
- x: torch.Tensor,
412
- ) -> torch.Tensor:
413
- """PyTorch-native implementation equivalent to forward().
414
-
415
- Refer to https://github.com/BryceZhuo/PolyCom?tab=readme-ov-file/README.md
416
- """
417
-
418
- orig_dtype = x.dtype
419
- x_float = x.to(torch.float32)
420
- output = (
421
- self.weight[0] * self._norm(x_float**3)
422
- + self.weight[1] * self._norm(x_float**2)
423
- + self.weight[2] * self._norm(x_float)
424
- + self.bias
425
- )
426
- return output.to(orig_dtype)
427
-
428
- def forward_cuda(
429
- self,
430
- x: torch.Tensor,
431
- ) -> torch.Tensor:
432
- return poly_norm(x, self.weight, self.bias, self.variance_epsilon)
433
-
434
-
372
class LayerNorm(nn.Module):
373
"""
374
Layer Normalization.