add gelt and rms_norm

Author: shihaobai

lightllm/models/vit/layer_infer/post_layer_infer.py

@@ -3,7 +3,7 @@
 import torch.distributed as dist
 from lightllm.models.vit.layer_weights.pre_and_post_layer_weight import ViTPreAndPostLayerWeight
 from lightllm.utils.dist_utils import get_current_rank_in_dp, get_dp_world_size
-
+from lightllm.models.vit.triton_kernel.gelu_vit import gelu
 
 class ViTPostLayerInfer:
     """ """
@@ -44,8 +44,9 @@ def forward(self, vit_embeds, layer_weight: ViTPreAndPostLayerWeight):
             layer_weight.mlp1_1_bias_, vit_embeds_norm.view(-1, vit_embeds_norm.shape[-1]), layer_weight.mlp1_1_weight_
         )
 
-        vit_embeds_gelu = torch.nn.functional.gelu(vit_embeds_1)
-
+        # vit_embeds_gelu = torch.nn.functional.gelu(vit_embeds_1)
+        vit_embeds_gelu = gelu(vit_embeds_1)
+        
         vit_embeds_out = torch.addmm(
             layer_weight.mlp1_3_bias_,
             vit_embeds_gelu.view(-1, self.llm_hidden_size // self.tp_world_size_),

lightllm/models/vit/layer_infer/transformer_layer_infer.py

@@ -10,7 +10,8 @@
 from lightllm.models.llama.triton_kernel.rmsnorm import rmsnorm_forward, torch_rms_norm
 from lightllm.models.vit.triton_kernel.flashattention_nopad import flash_attention_fwd
 from lightllm.utils.dist_utils import get_current_rank_in_dp, get_dp_world_size
-
+from lightllm.models.vit.triton_kernel.gelu_vit import gelu
+from lightllm.models.vit.triton_kernel.rms_norm_vit import rms_norm
 
 class ViTTransformerLayerInfer:
     """ """
@@ -58,7 +59,7 @@ def tp_norm(self, input, weight):
 
     def _att_norm(self, input, layer_weight: ViTTransformerLayerWeight) -> torch.Tensor:
         if layer_weight.norm_type == "rms_norm":
-            b = rmsnorm_forward(input, weight=layer_weight.att_norm_weight_.weight, eps=self.eps_)
+            b = rms_norm(input, weight=layer_weight.att_norm_weight_.weight, eps=self.eps_)
         else:
             b = torch.nn.functional.layer_norm(
                 input,
@@ -71,7 +72,7 @@ def _att_norm(self, input, layer_weight: ViTTransformerLayerWeight) -> torch.Ten
 
     def _ffn_norm(self, input, layer_weight: ViTTransformerLayerWeight) -> torch.Tensor:
         if layer_weight.norm_type == "rms_norm":
-            return rmsnorm_forward(input, weight=layer_weight.ffn_norm_weight_.weight, eps=self.eps_)
+            return rms_norm(input, weight=layer_weight.ffn_norm_weight_.weight, eps=self.eps_)
         else:
             return torch.nn.functional.layer_norm(
                 input,
@@ -113,7 +114,8 @@ def _get_o(self, input, layer_weight: ViTTransformerLayerWeight) -> torch.Tensor
 
     def _ffn(self, input, layer_weight: ViTTransformerLayerWeight) -> torch.Tensor:
         fc1 = layer_weight.ffn_1_proj_.mm(input.view(-1, self.embed_dim_), use_custom_tensor_mananger=False)
-        ffn1_out = torch.nn.functional.gelu(fc1)
+        # ffn1_out = torch.nn.functional.gelu(fc1)
+        ffn1_out = gelu(fc1)
         input_shape = input.shape
         input = None
         ffn2_out = layer_weight.ffn_2_proj_.mm(ffn1_out, use_custom_tensor_mananger=False)

lightllm/models/vit/triton_kernel/gelu_vit.py

@@ -0,0 +1,34 @@
+import torch
+import triton
+import triton.language as tl
+
+
+
+@triton.jit
+def gelu(x):
+    x_fp32 = x.to(tl.float32)
+    x_gelu = 0.5 * x_fp32 * (1 + tl.math.erf(x_fp32 * 0.7071067811))
+    return x_gelu
+
+# 定义 Triton 内核
+@triton.jit
+def gelu_kernel(output_ptr, input_ptr, n_elements, BLOCK_SIZE: tl.constexpr):
+    pid = tl.program_id(axis=0)
+    block_start = pid * BLOCK_SIZE
+    offsets = block_start + tl.arange(0, BLOCK_SIZE)
+    mask = offsets < n_elements
+    input = tl.load(input_ptr + offsets, mask=mask)
+    output = gelu(input)
+    tl.store(output_ptr + offsets, output, mask=mask)
+
+# 自定义 torch.autograd.Function
+class GeluTriton(torch.autograd.Function):
+    @staticmethod
+    def forward(ctx, input):
+        output = torch.empty_like(input)
+        n_elements = input.numel()
+        grid = lambda meta: (triton.cdiv(n_elements, meta['BLOCK_SIZE']),)
+        gelu_kernel[grid](output, input, n_elements, BLOCK_SIZE=1024)
+        return output
+
+gelu = GeluTriton.apply

lightllm/models/vit/triton_kernel/rms_norm_vit.py

@@ -0,0 +1,57 @@
+import torch
+import triton
+import triton.language as tl
+from torch import Tensor
+
+
+@triton.jit
+def rms_norm_kernel(
+    input,
+    weight,
+    output,
+    input_row_stride: tl.constexpr,
+    eps: tl.constexpr,
+    N_COLS: tl.constexpr,
+    BLOCK_N: tl.constexpr,
+):
+    """Rms norm kernel."""
+    prog_id = tl.program_id(0)
+    offsets = tl.arange(0, BLOCK_N)
+
+    w = tl.load(weight + offsets, mask=offsets < N_COLS)
+
+    x_ptr = input + prog_id * input_row_stride
+    x = tl.load(x_ptr + offsets, mask=offsets < N_COLS)
+    xf = x.to(tl.float32)
+
+    var = tl.sum(xf * xf, 0) * float(1.0 / N_COLS)
+    out = xf / tl.sqrt(var + eps)
+    out = (w * out).to(x.dtype)
+
+    out_ptr = output + prog_id * input_row_stride
+    tl.store(out_ptr + offsets, out, mask=offsets < N_COLS)
+
+
+def rms_norm(hidden_states: Tensor, weight: Tensor, eps: float = 1e-5):
+    """Rms norm."""
+    feat_size = weight.shape[0]
+    seq_len = hidden_states.numel() // hidden_states.size(-1)
+    input_stride = hidden_states.stride(-2)
+
+    BLOCK_N = triton.next_power_of_2(feat_size)
+    out = torch.empty_like(hidden_states)
+
+    grid = (seq_len,)
+    rms_norm_kernel[grid](
+        hidden_states,
+        weight,
+        out,
+        input_row_stride=input_stride,
+        eps=eps,
+        N_COLS=feat_size,
+        BLOCK_N=BLOCK_N,
+        num_warps=4,
+        num_stages=3,
+    )
+
+    return out