Sync v2.0 version of code to github repo

fastdeploy/model_executor/layers/attention/attention.py

@@ -14,12 +14,17 @@
 # limitations under the License.
 """
 
-from typing import Optional
+from typing import Dict, Optional
 
+import numpy as np
 import paddle
 from paddle import nn
+from paddleformers.utils.log import logger
 
-from fastdeploy.worker.model_runner import ForwardMeta
+from fastdeploy.config import FDConfig
+from fastdeploy.model_executor.layers.quantization.quant_base import \
+    QuantMethodBase
+from fastdeploy.worker.forward_meta import ForwardMeta
 
 
 class Attention(nn.Layer):
@@ -29,26 +34,24 @@ class Attention(nn.Layer):
 
     def __init__(
         self,
-        llm_config,
+        fd_config: FDConfig,
         layer_id: int,
-        logit_cap: float = 0.0,
         v_head_dim: int = -1,
         rope_type: str = "",
         qkv_bias: Optional[paddle.Tensor] = None,
         qkv_scale: Optional[paddle.Tensor] = None,
         prefix: str = "",
-        out_scale: float = -1.,
-        linear_shift=None,
-        linear_smooth=None,
-        use_neox_rotary_style=False,
+        out_scale: float = -1.0,
+        linear_shift: paddle.Tensor = None,
+        linear_smooth: paddle.Tensor = None,
+        use_neox_rotary_style: bool = False,
     ) -> None:
         """
         Initializes `LMLayer` with the given parameters.
 
         Args:
-            llm_config (dict): The config of LM model.
+            fd_config (dict): The config of LM model.
             layer_id (int): The id of current layer.
-            logit_cap (float, optional): The cap for logits. Defaults to 0.0.
             v_head_dim (int, optional): The head dim of value. Defaults to -1.
             rope_type (str, optional): The type of RoPE. Defaults to "".
             qkv_bias (Optional[paddle.Tensor], optional): The bias of QKV. Defaults to None.
@@ -61,34 +64,46 @@ class Attention(nn.Layer):
             ValueError: If the `v_head_dim` is less than 0.
         """
         super().__init__()
-        self.num_heads = llm_config.model_config.num_attention_heads // llm_config.parallel_config.mp_size
-        self.head_dim = llm_config.model_config.hidden_size // llm_config.model_config.num_attention_heads
-        self.kv_num_heads = llm_config.model_config.num_key_value_heads // llm_config.parallel_config.mp_size
-        self.layer_id = layer_id
-        self.logit_cap = logit_cap
-        self.v_head_dim = v_head_dim if v_head_dim > 0 else self.head_dim
-        self.rope_type = rope_type
-        self.qk_head_dim = self.head_dim
+        self.num_heads: int = fd_config.model_config.num_attention_heads // fd_config.parallel_config.tensor_parallel_degree
+        self.head_dim: int = fd_config.model_config.head_dim
+        self.kv_num_heads: int = \
+            fd_config.model_config.num_key_value_heads // fd_config.parallel_config.tensor_parallel_degree
+        self.layer_id: int = layer_id
+        self.v_head_dim: int = v_head_dim if v_head_dim > 0 else self.head_dim
+        self.rope_type: str = rope_type
+        self.qk_head_dim: int = self.head_dim
+        self.prefix: str = prefix
         # not use
-        self.tp_q_head_num = self.num_heads
-        self.tp_k_head_num = self.num_heads
-        self.tp_v_head_num = self.num_heads
-        # not use
-        self.scaling = 1.0 / (self.head_dim**0.5)
-        self.linear_shift = linear_shift
-        self.linear_smooth = linear_smooth
-        self.qkv_bias = qkv_bias
-        self.qkv_scale = qkv_scale
+        self.linear_shift: paddle.Tensor | None = linear_shift
+        self.linear_smooth: paddle.Tensor | None = linear_smooth
+        self.qkv_bias: paddle.Tensor | None = qkv_bias
+        self.qkv_scale: paddle.Tensor | None = qkv_scale
         self._dtype = self._helper.get_default_dtype()
-        self.out_scale = out_scale
-        self.use_neox_rotary_style = use_neox_rotary_style
-        if llm_config.kvcache_config is not None:
-            self.kvcache_quant_method = llm_config.kvcache_config.kvcache_quant_config.get_quant_method(
+
+        self.out_scale: float = out_scale
+        self.use_neox_rotary_style: bool = use_neox_rotary_style
+
+        if fd_config.quant_config and hasattr(fd_config.quant_config,
+                                              "kv_cache_quant_type"):
+            self.kvcache_quant_method: QuantMethodBase = fd_config.quant_config.get_quant_method(
                 self)
-            self.kvcache_quant_method.create_weights(self)
-        if llm_config.quant_config is not None:
-            self.quant_max_bound = llm_config.quant_config.quant_max_bound
-            self.quant_min_bound = llm_config.quant_config.quant_min_bound
+        else:
+            self.kvcache_quant_method = None
+
+        if self.kvcache_quant_method is None:
+            logger.info(f"Attention is running in cache kv {self._dtype} mode")
+        else:
+            logger.info(
+                f"Attention is running in cache kv {self.kvcache_quant_method.cache_quant_config.quant_type} mode"
+            )
+
+    def load_state_dict(self, state_dict: Dict[str,
+                                               paddle.Tensor | np.ndarray]):
+        '''
+        Attention only have quant related scales not other parameters.
+        '''
+        if self.kvcache_quant_method is not None:
+            self.kvcache_quant_method.create_weights(self, state_dict)
 
     def forward(
         self,
@@ -97,7 +112,7 @@ class Attention(nn.Layer):
         v: paddle.Tensor = None,
         qkv: paddle.Tensor = None,
         forward_meta: ForwardMeta = None,
-    ):
+    ) -> paddle.Tensor:
         """
         The forward function of attention layer.
         args: