[bugfix][#280] MLA layer size calculated wrong (#281)

* Fix the layer_size calculated wrong for mla

* Fix the style

Author: maobaolong

ucm/integration/vllm/uc_connector.py

@@ -187,9 +187,9 @@ def DataOffset(self, kv_layer, rank, layer_id, is_v):
             kv_layer[1][0].numel() if not self.is_mla else 0
         ) * elem_size
         # When tp > 1 layer_size = (k_min_data_block_size + v_min_data_block_size) * tp_size
-        layer_size = (
-            k_min_data_block_size + v_min_data_block_size
-        ) * self.total_tp_size
+        layer_size = (k_min_data_block_size + v_min_data_block_size) * (
+            self.total_tp_size if not self.is_mla else 1
+        )
         if is_v:
             # Offset of v = Offset of k + k_min_data_block_size
             return int(