Compile layernorm

Signed-off-by: Chang Liu <9713593+chang-l@users.noreply.github.com>

Author: chang-l

tensorrt_llm/_torch/modules/layer_norm.py

@@ -18,6 +18,8 @@
 import torch
 from torch import nn
 
+from ..utils import maybe_compile
+
 
 class LayerNorm(nn.Module):
     """Layer normalization module with configurable weight and bias parameters.
@@ -65,6 +67,7 @@ def __init__(
                                  persistent=False)
         self.variance_epsilon = eps
 
+    @maybe_compile(dynamic=True)
     def forward(
         self,
         hidden_states: torch.Tensor,