change layernorm kernel

Author: flyinglandlord

lightllm/models/cohere/layer_infer/post_layer_infer.py

@@ -22,10 +22,9 @@ def __init__(self, tp_rank, world_size, network_config, mode):
         return
 
     def _norm(self, input, infer_state, layer_weight: CoherePreAndPostLayerWeight) -> torch.Tensor:
-        input = layernorm_forward(
+        return layernorm_forward(
             input.unsqueeze(1), layer_weight.final_norm_weight_.unsqueeze(0), eps=self.eps_
-        ).squeeze_(1)
-        return input
+        ).squeeze(1)
 
     def _slice_get_last_input(self, input_embdings, infer_state: CohereInferStateInfo):
         if infer_state.is_splitfuse:

lightllm/models/cohere/layer_infer/transformer_layer_infer.py

@@ -6,7 +6,7 @@
 )
 from lightllm.models.cohere.infer_struct import CohereInferStateInfo
 from lightllm.models.cohere.layer_weights.transformer_layer_weight import CohereTransformerLayerWeight
-from lightllm.models.cohere.triton_kernels.layernorm import layernorm_forward
+from lightllm.models.cohere.triton_kernels.layernorm import layernorm_forward, torch_layernorm
 from lightllm.models.cohere.triton_kernels.rotary_emb import rotary_emb_fwd
 from lightllm.models.llama.layer_infer.transformer_layer_infer import LlamaTransformerLayerInfer
 from lightllm.models.llama.triton_kernel.silu_and_mul import silu_and_mul_fwd
@@ -42,16 +42,13 @@ def _bind_rotary_emb_fwd(self):
         self._rotary_emb_fwd = partial(CohereTransformerLayerInfer._rotary_emb_fwd, self)
 
     def _att_norm(self, input, infer_state, layer_weight):
-        input = layernorm_forward(input.unsqueeze(1), layer_weight.att_norm_weight_.unsqueeze(0), self.eps_).squeeze_(1)
-        return input
+        return layernorm_forward(input.unsqueeze(1), layer_weight.att_norm_weight_.unsqueeze(0), self.eps_).squeeze(1)
 
     def _q_norm(self, input, infer_state, layer_weight):
-        input = layernorm_forward(input, layer_weight.q_norm_weight_, self.eps_)
-        return input
+        return layernorm_forward(input, layer_weight.q_norm_weight_, self.eps_)
 
     def _k_norm(self, input, infer_state, layer_weight):
-        input = layernorm_forward(input, layer_weight.k_norm_weight_, self.eps_)
-        return input
+        return layernorm_forward(input, layer_weight.k_norm_weight_, self.eps_)
 
     def _bind_norm(self):
         self._att_norm = partial(CohereTransformerLayerInfer._att_norm, self)

lightllm/models/cohere/model.py

@@ -4,6 +4,7 @@
 from lightllm.common.basemodel.layer_infer.template.transformer_layer_infer_cohere_template import (
     TransformerLayerCohereInferTpl,
 )
+from lightllm.common.mem_manager import MemoryManager
 from lightllm.models.cohere.infer_struct import CohereInferStateInfo
 from lightllm.models.cohere.layer_infer.post_layer_infer import CoherePostLayerInfer
 from lightllm.models.cohere.layer_infer.transformer_layer_infer import CohereTransformerLayerInfer

lightllm/models/cohere/triton_kernels/layernorm.py

@@ -13,6 +13,9 @@ def _layer_norm_fwd_kernel(
     stride_x_N,
     stride_x_hn,
     stride_x_hd,
+    stride_y_N,
+    stride_y_hn,
+    stride_y_hd,
     stride_w_hn,
     stride_w_hd,
     N,  # number of columns in X
@@ -23,7 +26,7 @@ def _layer_norm_fwd_kernel(
     H = tl.program_id(1)
 
     X += Seq * stride_x_N + H * stride_x_hn
-    Y += Seq * stride_x_N + H * stride_x_hn
+    Y += Seq * stride_y_N + H * stride_y_hn
     W += H * stride_w_hn
 
     _mean = tl.zeros([BLOCK_SIZE], dtype=tl.float32)
@@ -63,32 +66,41 @@ def layernorm_forward(
     assert X.shape[-1] == W.shape[-1]
     assert X.shape[-2] == W.shape[-2]
 
+    y = torch.empty_like(X)
+
     stride_x_N = X.stride(0)
     stride_x_hn = X.stride(1)
     stride_x_hd = X.stride(2)
+
+    stride_y_N = y.stride(0)
+    stride_y_hn = y.stride(1)
+    stride_y_hd = y.stride(2)
+
     stride_w_hn = W.stride(0)
     stride_w_hd = W.stride(1)
+
     N = X.shape[-1]
     BLOCK_SIZE = 128
 
-    Y = torch.empty_like(X)
-
     grid = (X.shape[0], X.shape[1])
     _layer_norm_fwd_kernel[grid](
         X,
         W,
-        Y,
+        y,
         stride_x_N,
         stride_x_hn,
         stride_x_hd,
+        stride_y_N,
+        stride_y_hn,
+        stride_y_hd,
         stride_w_hn,
         stride_w_hd,
         N,
         eps,
         BLOCK_SIZE,
     )
 
-    return Y
+    return y
 
 
 def torch_layernorm(x, weight, eps):
@@ -104,12 +116,12 @@ def torch_layernorm(x, weight, eps):
 def test_layernorm(eps=1e-5):
     # create data
     dtype = torch.float16
-    x_shape = (1000, 1, 128)
+    x_shape = (5, 1, 128)
     w_shape = (x_shape[-2], x_shape[-1])
     weight = torch.rand(w_shape, dtype=dtype, device="cuda")
     x = -2.3 + 0.5 * torch.randn(x_shape, dtype=dtype, device="cuda")
     # forward pass
-    y_ref = torch_layernorm(x.to(torch.float32), weight.to(torch.float32), eps).to(dtype)
+    y_ref = torch_layernorm(x, weight, eps).to(dtype)
     y_out = layernorm_forward(x, weight, eps)
 
     # compare