Initializes a KV batch linear layer that internally splits into K and V projections. Args: fd_config (FDConfig): Inference-related parameters. prefix (str): Unique name of the layer, used to name internal attributes. kv_lora_rank (int): LoRA rank
(
self,
fd_config: FDConfig,
kv_b_proj: nn.Layer,
prefix: str = "",
kv_lora_rank: int = None,
num_attention_heads: int = None,
qk_nope_head_dim: int = None,
v_head_dim: int = None,
)
| 946 | """ |
| 947 | |
| 948 | def __init__( |
| 949 | self, |
| 950 | fd_config: FDConfig, |
| 951 | kv_b_proj: nn.Layer, |
| 952 | prefix: str = "", |
| 953 | kv_lora_rank: int = None, |
| 954 | num_attention_heads: int = None, |
| 955 | qk_nope_head_dim: int = None, |
| 956 | v_head_dim: int = None, |
| 957 | ): |
| 958 | """ |
| 959 | Initializes a KV batch linear layer that internally splits into K and V projections. |
| 960 | |
| 961 | Args: |
| 962 | fd_config (FDConfig): Inference-related parameters. |
| 963 | prefix (str): Unique name of the layer, used to name internal attributes. |
| 964 | kv_lora_rank (int): LoRA rank for KV projection. Defaults to None. |
| 965 | num_attention_heads (int): Number of attention heads. Defaults to None. |
| 966 | qk_nope_head_dim (int): Dimension for Q/K projection (nope part). Defaults to None. |
| 967 | v_head_dim (int): Dimension for V projection. Defaults to None. |
| 968 | with_bias (bool): Whether to include bias or not. Defaults to False. |
| 969 | """ |
| 970 | super().__init__() |
| 971 | self.tp_size = fd_config.parallel_config.tensor_parallel_size |
| 972 | self.kv_lora_rank = kv_lora_rank |
| 973 | self.num_attention_heads = num_attention_heads |
| 974 | self.qk_nope_head_dim = qk_nope_head_dim |
| 975 | self.v_head_dim = v_head_dim |
| 976 | # Split num_attention_heads when using TP inference. |
| 977 | self.num_heads_per_partition = divide(num_attention_heads, self.tp_size) |
| 978 | self.local_rank = fd_config.parallel_config.tensor_parallel_rank |
| 979 | self.fd_config = fd_config |
| 980 | if self.fd_config.load_config.load_choices == "default_v1": |
| 981 | self.kv_b_proj = kv_b_proj |
| 982 | else: |
| 983 | self.kv_b_proj = None |
| 984 | |
| 985 | self.weight_dtype = self._helper.get_default_dtype() |
| 986 | |
| 987 | # Override weight keys to use the combined kv_b_proj |
| 988 | self.weight_key = f"{prefix}.weight" # e.g., "kv_b_proj.weight" |
| 989 | |
| 990 | def process_weights_after_loading(self): |
| 991 | if self.fd_config.load_config.dynamic_load_weight: |
nothing calls this directly
no test coverage detected