fp8 scale repeat for qwen3 (#879)

Co-authored-by: baishihao <[email protected]>
Co-authored-by: wangzaijun <[email protected]>

Author: shihaobai

lightllm/models/qwen3_moe/layer_weights/transformer_layer_weight.py

@@ -60,15 +60,21 @@ def _repeat_weight(self, name, weights):
         if name in weights:
             weights[name] = (
                 weights[name]
-                .reshape(self.network_config_["num_key_value_heads"], self.head_dim, -1)
+                .reshape(self.network_config_["num_key_value_heads"], -1, weights[name].shape[1])
                 .unsqueeze(1)
                 .repeat(repeat_params)
-                .reshape(self.network_config_["num_key_value_heads"] * self.head_dim * repeat_size, -1)
+                .reshape(-1, weights[name].shape[1])
             )
 
     def load_hf_weights(self, weights):
         self._repeat_weight(self._k_weight_name, weights)
         self._repeat_weight(self._v_weight_name, weights)
+        kv_b_quant_method = self.quant_cfg.get_quant_method(self.layer_num_, "kv_b_proj")
+        if self.quant_cfg.quantized_weight:
+            _k_scale_weight_name = self._k_weight_name.replace("weight", kv_b_quant_method.weight_scale_suffix)
+            self._repeat_weight(_k_scale_weight_name, weights)
+            _v_scale_weight_name = self._v_weight_name.replace("weight", kv_b_quant_method.weight_scale_suffix)
+            self._repeat_weight(_v_scale_weight_name, weights)
         return super().load_hf_weights(weights)
 
     def _init_weight(self):

lightllm/models/qwen3_moe/model.py

@@ -26,6 +26,16 @@ def _verify_params(self):
         assert self.config["num_attention_heads"] % self.tp_world_size_ == 0
         return
 
+    def _init_some_value(self):
+        # Dealing with head_dim_!=n_embed // num_attention_heads scenarios, such as mistral 13B
+        head_dim_ = self.config["n_embed"] // self.config["num_attention_heads"]
+        self.head_dim_ = self.config.get("head_dim", head_dim_)
+        self.tp_k_head_num_ = max(self.config["num_key_value_heads"] // self.tp_world_size_, 1)
+        self.tp_v_head_num_ = self.tp_k_head_num_
+        self.layers_num = self.config["n_layer"]
+        self.vocab_size = self.config["vocab_size"]
+        return
+
     def _init_mem_manager(self):
         head_dim_ = self.config["hidden_size"] // self.config["num_attention_heads"]
         head_dim_ = self.config.get("head_dim", head_dim_)