init

vllm_vacc/vllm/model_executor/layers/layernorm.py

@@ -0,0 +1,32 @@
+from typing import Optional, Tuple, Union
+import torch
+
+
+def RMSNorm_forward_vacc(
+    self,
+    x: torch.Tensor,
+    residual: Optional[torch.Tensor] = None,
+) -> Union[torch.Tensor, Tuple[torch.Tensor, torch.Tensor]]:
+    # if residual is not None:
+    #     x = x + residual
+    #     residual = x
+    hidden_size = x.shape[-1]
+    if hidden_size != self.hidden_size:
+        raise ValueError("Expected hidden_size to be "
+                         f"{self.hidden_size}, but found: {hidden_size}")
+    if self.variance_size_override is None:
+        x_var = x
+    else:
+        if hidden_size < self.variance_size_override:
+            raise ValueError(
+                "Expected hidden_size to be at least "
+                f"{self.variance_size_override}, but found: {hidden_size}")
+        x_var = x[:, :, :self.variance_size_override]
+    # x_var=x_var.unsqueeze(0)
+    # out = torch.vacc.rms_norm(x_var,self.weight,self.variance_epsilon)
+    # if residual is None:
+    #     return out.squeeze(0)
+    # else:
+    #     return out.squeeze(0), residual
+    out = torch.vacc.fused_residual_rmsnorm(x_var, self.weight, residual, self.variance_epsilon, x_var, residual)
+    return out