fix rmsnorm

Author: sangchengmeng

lightllm/models/deepseek2/layer_infer/transformer_layer_infer.py

@@ -154,18 +154,16 @@ def _get_qkv(
             q = layer_weight.q_weight_.mm(input)
         else:
             q = layer_weight.q_a_proj_.mm(input)
-            q = rmsnorm_forward(
-                q, weight=layer_weight.q_a_layernorm_.weight, eps=self.eps_, use_custom_tensor_mananger=True
-            )
+            rmsnorm_forward(q, weight=layer_weight.q_a_layernorm_.weight, eps=self.eps_, out=q)
             q = layer_weight.q_b_proj_.mm(q)
         q = q.view(-1, self.tp_q_head_num_, self.qk_nope_head_dim + self.qk_rope_head_dim)
         q_nope, q_rope = torch.split(q, [self.qk_nope_head_dim, self.qk_rope_head_dim], dim=-1)
         layer_weight.kv_a_proj_with_mqa_.mm(input, out=cache_kv.view(-1, self.kv_lora_rank + self.qk_rope_head_dim))
-        cache_kv[:, :, : self.kv_lora_rank] = rmsnorm_forward(
+        rmsnorm_forward(
             cache_kv[:, :, : self.kv_lora_rank],
             weight=layer_weight.kv_a_layernorm_.weight,
             eps=self.eps_,
-            use_custom_tensor_mananger=True,
+            out=cache_kv[:, :, : self.kv_lora_rank],
         )
 
         rotary_emb_fwd(
@@ -193,16 +191,16 @@ def _tpsp_get_qkv(
             q = layer_weight.q_weight_.mm(input)
         else:
             q = layer_weight.q_a_proj_.mm(input)
-            q = rmsnorm_forward(q, weight=layer_weight.q_a_layernorm_.weight, eps=self.eps_)
+            rmsnorm_forward(q, weight=layer_weight.q_a_layernorm_.weight, eps=self.eps_, out=q)
             q = layer_weight.q_b_proj_.mm(q)
         q = q.view(-1, self.tp_q_head_num_, self.qk_nope_head_dim + self.qk_rope_head_dim)
         q_nope, q_rope = torch.split(q, [self.qk_nope_head_dim, self.qk_rope_head_dim], dim=-1)
         layer_weight.kv_a_proj_with_mqa_.mm(input, out=cache_kv.view(-1, self.kv_lora_rank + self.qk_rope_head_dim))
-        cache_kv[:, :, : self.kv_lora_rank] = rmsnorm_forward(
+        rmsnorm_forward(
             cache_kv[:, :, : self.kv_lora_rank],
             weight=layer_weight.kv_a_layernorm_.weight,
             eps=self.eps_,
-            use_custom_tensor_mananger=True,
+            out=cache_kv[:, :, : self.kv_lora_rank],
         )
         rotary_emb_fwd(
             q_rope,

lightllm/models/deepseek_mtp/layer_infer/pre_layer_infer.py

@@ -20,12 +20,8 @@ def _mtp_context_forward(
     ):
         tgt_embdings = infer_state.deepseekv3_mtp_draft_input_hiddens
         assert input_embdings.shape[0] == tgt_embdings.shape[0]
-        input_embdings = rmsnorm_forward(
-            input_embdings, weight=layer_weight.enorm_weight_, eps=self.eps_, use_custom_tensor_mananger=True
-        )
-        tgt_embdings = rmsnorm_forward(
-            tgt_embdings, weight=layer_weight.hnorm_weight_, eps=self.eps_, use_custom_tensor_mananger=True
-        )
+        rmsnorm_forward(input_embdings, weight=layer_weight.enorm_weight_, eps=self.eps_, out=input_embdings)
+        rmsnorm_forward(tgt_embdings, weight=layer_weight.hnorm_weight_, eps=self.eps_, out=tgt_embdings)
 
         cat_embdings = torch.cat((input_embdings, tgt_embdings), dim=-1)
 
@@ -40,12 +36,8 @@ def _mtp_token_forward(
     ):
         tgt_embdings = infer_state.deepseekv3_mtp_draft_input_hiddens
         assert input_embdings.shape[0] == tgt_embdings.shape[0]
-        input_embdings = rmsnorm_forward(
-            input_embdings, weight=layer_weight.enorm_weight_, eps=self.eps_, use_custom_tensor_mananger=True
-        )
-        tgt_embdings = rmsnorm_forward(
-            tgt_embdings, weight=layer_weight.hnorm_weight_, eps=self.eps_, use_custom_tensor_mananger=True
-        )
+        rmsnorm_forward(input_embdings, weight=layer_weight.enorm_weight_, eps=self.eps_, out=input_embdings)
+        rmsnorm_forward(tgt_embdings, weight=layer_weight.hnorm_weight_, eps=self.eps_, out=tgt_embdings)
 
         cat_embdings = torch.cat((input_embdings, tgt_embdings), dim=-1)
 

lightllm/models/llama/layer_infer/post_layer_infer.py

@@ -25,7 +25,7 @@ def __init__(self, network_config, mode):
         return
 
     def _norm(self, input, infer_state, layer_weight: LlamaPreAndPostLayerWeight) -> torch.Tensor:
-        return rmsnorm_forward(input, layer_weight.final_norm_weight_, eps=self.eps_, use_custom_tensor_mananger=True)
+        return rmsnorm_forward(input, layer_weight.final_norm_weight_, eps=self.eps_)
 
     def _slice_get_last_input(self, input_embdings, infer_state: LlamaInferStateInfo):
 

lightllm/models/llama/layer_infer/transformer_layer_infer.py

@@ -135,16 +135,16 @@ def _bind_attention(self):
     def _att_norm(
         self, input, infer_state: LlamaInferStateInfo, layer_weight: LlamaTransformerLayerWeight
     ) -> torch.Tensor:
-        return rmsnorm_forward(
-            input, weight=layer_weight.att_norm_weight_.weight, eps=self.eps_, use_custom_tensor_mananger=True
-        )
+        out = self.alloc_tensor(input.shape, input.dtype)
+        rmsnorm_forward(input, weight=layer_weight.att_norm_weight_.weight, eps=self.eps_, out=out)
+        return out
 
     def _ffn_norm(
         self, input, infer_state: LlamaInferStateInfo, layer_weight: LlamaTransformerLayerWeight
     ) -> torch.Tensor:
-        return rmsnorm_forward(
-            input, weight=layer_weight.ffn_norm_weight_.weight, eps=self.eps_, use_custom_tensor_mananger=True
-        )
+        out = self.alloc_tensor(input.shape, input.dtype)
+        rmsnorm_forward(input, weight=layer_weight.ffn_norm_weight_.weight, eps=self.eps_, out=out)
+        return out
 
     def _get_qkv(
         self, input, cache_kv, infer_state: LlamaInferStateInfo, layer_weight: LlamaTransformerLayerWeight

lightllm/models/llama/triton_kernel/rmsnorm.py

@@ -2,11 +2,10 @@
 import torch
 import triton
 import triton.language as tl
-from lightllm.common.basemodel.layer_infer.cache_tensor_manager import g_cache_manager
 
 
 @triton.jit
-def _rms_norm_low_accuracy_kernel(
+def _rmsnorm_kernel(
     X,  # pointer to the input
     Y,  # pointer to the output
     W,  # pointer to the weights
@@ -42,15 +41,9 @@ def _rms_norm_low_accuracy_kernel(
         tl.store(Y + cols * y_stride1, y.to(Y.dtype.element_ty), mask=mask)
 
 
-def rmsnorm_forward_low_accuracy(x: torch.Tensor, weight, eps, use_custom_tensor_mananger: bool = False):
+def rmsnorm(x: torch.Tensor, weight, eps, out=None):
     # allocate output
-    if use_custom_tensor_mananger:
-        shape = x.shape
-        dtype = x.dtype
-        device = x.device
-        y = g_cache_manager.alloc_tensor(shape, dtype, device=device)
-    else:
-        y = torch.empty_like(x)
+    y = torch.empty_like(x) if out is None else out
     # reshape input data into 2D tensor
     x_arg = x.view(-1, x.shape[-1])
     y_arg = y.view(-1, x.shape[-1])
@@ -68,7 +61,7 @@ def rmsnorm_forward_low_accuracy(x: torch.Tensor, weight, eps, use_custom_tensor
     if BLOCK_SIZE > 16384:
         BLOCK_SIZE = 16384
     # enqueue kernel
-    _rms_norm_low_accuracy_kernel[(M,)](
+    _rmsnorm_kernel[(M,)](
         x_arg,
         y_arg,
         weight,
@@ -85,7 +78,7 @@ def rmsnorm_forward_low_accuracy(x: torch.Tensor, weight, eps, use_custom_tensor
 
 
 @triton.jit
-def _rms_norm_high_accuracy_kernel(
+def _rms_norm_kernel(
     input,
     weight,
     output,
@@ -115,9 +108,7 @@ def _rms_norm_high_accuracy_kernel(
     tl.store(out_ptr + offsets * out_col_stride, out, mask=offsets < N_COLS)
 
 
-def rmsnorm_forward_high_accuracy(
-    hidden_states: torch.Tensor, weight: torch.Tensor, eps: float = 1e-5, use_custom_tensor_mananger: bool = False
-):
+def rms_norm(hidden_states: torch.Tensor, weight: torch.Tensor, eps: float = 1e-5, out=None):
     """Rms norm."""
 
     assert hidden_states.is_contiguous(), "hidden_states must be contiguous"
@@ -133,17 +124,12 @@ def rmsnorm_forward_high_accuracy(
     in_row_stride, in_col_stride = hidden_states.stride(0), hidden_states.stride(1)
 
     BLOCK_N = triton.next_power_of_2(hidden_dim)
-    if use_custom_tensor_mananger:
-        shape = hidden_states.shape
-        dtype = hidden_states.dtype
-        device = hidden_states.device
-        output = g_cache_manager.alloc_tensor(shape, dtype, device=device)
-    else:
-        output = torch.empty_like(hidden_states)
+
+    output = torch.empty_like(hidden_states) if out is None else out
 
     out_row_stride, out_col_stride = output.stride(0), output.stride(1)
     grid = (rows,)
-    _rms_norm_high_accuracy_kernel[grid](
+    _rms_norm_kernel[grid](
         hidden_states,
         weight,
         output,
@@ -171,21 +157,21 @@ def test_rms_norm(M, N, dtype, eps=1e-5, device="cuda"):
     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_tri = rmsnorm_forward_low_accuracy(x, weight, eps)
-    y_tri_high_acc = rmsnorm_forward_high_accuracy(x, weight, eps)
+    y_tri = rmsnorm_forward(x, weight, eps)
+    y_tri_1 = rms_norm(x, weight, eps)
     y_ref = torch_rms_norm(x.to(torch.float32), weight.to(torch.float32), eps).to(dtype)
 
     # compare
-    print("type:", y_tri.dtype, y_ref.dtype, y_tri_high_acc.dtype)
+    print("type:", y_tri.dtype, y_ref.dtype, y_tri_1.dtype)
     print("max delta:", torch.max(torch.abs(y_tri - y_ref)))
-    print("max delta:", torch.max(torch.abs(y_tri_high_acc - y_ref)))
+    print("max delta:", torch.max(torch.abs(y_tri_1 - y_ref)))
     assert torch.allclose(y_tri, y_ref, atol=1e-2, rtol=0)
     return
 
 
 use_high_acc = os.getenv("RMSNORM_HIGH_ACCURACY", "False").upper() in ["ON", "TRUE", "1"]
 
 if use_high_acc:
-    rmsnorm_forward = rmsnorm_forward_high_accuracy
+    rmsnorm_forward = rms_norm
 else:
-    rmsnorm_forward = rmsnorm_forward_low_accuracy
+    rmsnorm_forward = rmsnorm

lightllm/models/qwen3/layer_infer/transformer_layer_infer.py

@@ -36,18 +36,18 @@ def _get_qkv(
             input, out=cache_kv.view(-1, (self.tp_k_head_num_ + self.tp_v_head_num_) * self.head_dim_)
         ).view(-1, (self.tp_k_head_num_ + self.tp_v_head_num_), self.head_dim_)
 
-        q = rmsnorm_forward(
+        rmsnorm_forward(
             q.view(-1, self.head_dim_),
             weight=layer_weight.q_norm_weight_.weight,
             eps=self.eps_,
             use_custom_tensor_mananger=True,
+            out=q.view(-1, self.head_dim_),
         )
 
         cache_kv[:, : self.tp_k_head_num_, :] = rmsnorm_forward(
             cache_kv[:, : self.tp_k_head_num_, :].reshape(-1, cache_kv.shape[-1]),
             weight=layer_weight.k_norm_weight_.weight,
             eps=self.eps_,
-            use_custom_tensor_mananger=True,
         ).view(-1, self.tp_k_head_num_, cache_kv.shape[-1])
 
         rotary_emb_fwd(

lightllm/models/qwen3_moe/layer_infer/transformer_layer_infer.py

@@ -60,18 +60,17 @@ def _get_qkv(
         cache_kv = layer_weight.kv_proj.mm(
             input, out=cache_kv.view(-1, (self.tp_k_head_num_ + self.tp_v_head_num_) * self.head_dim_)
         ).view(-1, (self.tp_k_head_num_ + self.tp_v_head_num_), self.head_dim_)
-        q = rmsnorm_forward(
+        rmsnorm_forward(
             q.view(-1, self.head_dim_),
             weight=layer_weight.q_norm_weight_.weight,
             eps=self.eps_,
-            use_custom_tensor_mananger=True,
+            out=q.view(-1, self.head_dim_),
         )
 
         cache_kv[:, : self.tp_k_head_num_, :] = rmsnorm_forward(
             cache_kv[:, : self.tp_k_head_num_, :].reshape(-1, cache_kv.shape[-1]),
             weight=layer_weight.k_norm_weight_.weight,
             eps=self.eps_,
-            use_custom_tensor_mananger=True,
         ).view(-1, self.tp_k_head_num_, cache_kv.shape[-1])
 
         rotary_emb_fwd(

lightllm/models/vit/layer_infer/transformer_layer_infer.py

@@ -10,7 +10,7 @@
 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_fwd
-from lightllm.models.llama.triton_kernel.rmsnorm import rmsnorm_forward_high_accuracy as rms_norm
+from lightllm.models.llama.triton_kernel.rmsnorm import rms_norm
 from lightllm.common.basemodel.layer_infer.cache_tensor_manager import g_cache_manager