add flashinfer fused_norm_add op

Author: sangchengmeng

lightllm/models/deepseek2/layer_infer/transformer_layer_infer.py

@@ -821,8 +821,8 @@ def overlap_tpsp_token_forward(
         _0_o = self._token_attention_kernel(_0_q, infer_state, layer_weight)
         _0_q = None
         _0_o = self._tpsp_get_o(_0_o, infer_state, layer_weight)
-        _0_input1 = _0_o.view(-1, self.embed_dim_)
         if FLASHINFER_AVAILABLE and layer_weight.norm_type == "rms_norm":
+            _0_input1 = _0_o.view(-1, self.embed_dim_)
             flashinfer.norm.fused_add_rmsnorm(
                 _0_input1, input_embdings, layer_weight.ffn_norm_weight_.weight, eps=self.eps_
             )
@@ -952,8 +952,8 @@ def overlap_tpsp_context_forward(
         _0_o = self._context_attention_kernel(_0_q, _0_cache_kv, infer_state, layer_weight)
         _0_q = None
         _0_o = self._tpsp_get_o(_0_o, infer_state, layer_weight)
-        _0_input1 = _0_o.view(-1, self.embed_dim_)
         if FLASHINFER_AVAILABLE and layer_weight.norm_type == "rms_norm":
+            _0_input1 = _0_o.view(-1, self.embed_dim_)
             flashinfer.norm.fused_add_rmsnorm(
                 _0_input1, input_embdings, layer_weight.ffn_norm_weight_.weight, eps=self.eps_
             )

lightllm/models/qwen3_moe/layer_infer/transformer_layer_infer.py

@@ -227,8 +227,8 @@ def overlap_tpsp_token_forward(
         _0_o = self._token_attention_kernel(_0_q, infer_state, layer_weight)
         _0_q = None
         _0_o = self._tpsp_get_o(_0_o, infer_state, layer_weight)
-        _0_input1 = _0_o.view(-1, self.embed_dim_)
         if FLASHINFER_AVAILABLE and layer_weight.norm_type == "rms_norm":
+            _0_input1 = _0_o.view(-1, self.embed_dim_)
             flashinfer.norm.fused_add_rmsnorm(
                 _0_input1, input_embdings, layer_weight.ffn_norm_weight_.weight, eps=self.eps_
             )
@@ -261,8 +261,8 @@ def overlap_tpsp_token_forward(
         _1_o = self._token_attention_kernel(_1_q, infer_state1, layer_weight)
         _1_q = None
         _1_o = self._tpsp_get_o(_1_o, infer_state1, layer_weight)
-        _1_input1 = _1_o.view(-1, self.embed_dim_)
         if FLASHINFER_AVAILABLE and layer_weight.norm_type == "rms_norm":
+            _1_input1 = _1_o.view(-1, self.embed_dim_)
             flashinfer.norm.fused_add_rmsnorm(
                 _1_input1, input_embdings1, weight=layer_weight.ffn_norm_weight_.weight, eps=self.eps_
             )
@@ -351,8 +351,8 @@ def overlap_tpsp_context_forward(
         _0_o = self._context_attention_kernel(_0_q, _0_cache_kv, infer_state, layer_weight)
         _0_q = None
         _0_o = self._tpsp_get_o(_0_o, infer_state, layer_weight)
-        _0_input1 = _0_o.view(-1, self.embed_dim_)
         if FLASHINFER_AVAILABLE and layer_weight.norm_type == "rms_norm":
+            _0_input1 = _0_o.view(-1, self.embed_dim_)
             flashinfer.norm.fused_add_rmsnorm(
                 _0_input1, input_embdings, weight=layer_weight.ffn_norm_weight_.weight, eps=self.eps_
             )