fused_moe_fp8接入lmslim

vllm/model_executor/layers/fused_moe/fused_moe.py

@@ -2058,7 +2058,7 @@ def fused_experts_impl(
                               max(N, K if not use_nn_moe else w2.shape[2]),
                               device=hidden_states.device,
                               dtype=hidden_states.dtype)
-    if use_int8_w8a8 is True:
+    if use_int8_w8a8 or use_fp8_w8a8:
         return fused_experts_impl_int8(hidden_states=hidden_states,
                                        w1=w1,
                                        w2=w2,
@@ -2068,8 +2068,8 @@ def fused_experts_impl(
                                        inplace=inplace,
                                        activation=activation,
                                        apply_router_weight_on_input=apply_router_weight_on_input,
-                                       use_fp8_w8a8=False,
-                                       use_int8_w8a8=True,
+                                       use_fp8_w8a8=use_fp8_w8a8,
+                                       use_int8_w8a8=use_int8_w8a8,
                                        use_int8_w8a16=False,
                                        use_int4_w4a16=False,
                                        per_channel_quant=per_channel_quant,

vllm/utils/__init__.py

@@ -2308,7 +2308,7 @@ class W8a8GetCacheJSON:
         self.moe_weight_shapes=[]
         arch_name = torch.cuda.get_device_properties("cuda").gcnArchName.split(':')[0]
         arch_cu = torch.cuda.get_device_properties(torch.cuda.current_device()).multi_processor_count
-        
+        self.cache_json_data = {}
         device_name =arch_name+'_'+str(arch_cu)+'cu'
         self.device_name=device_name
         self.topk=1
@@ -2410,20 +2410,27 @@ class W8a8GetCacheJSON:
         return self.triton_json_dir+f"/linear_{n}_{k}_block[{block_n},{block_k}]_{self.device_name}.json"
 
     def get_moeint8json_name(self,E,N1,N2,K,TOPK,
-                             block_size:Optional[list]=None,use_int4_w4a8:Optional[bool]=False):
+                             block_size:Optional[list]=None,use_int4_w4a8:Optional[bool]=False,use_int8_w8a8:Optional[bool]=False):
         if use_int4_w4a8:
             if block_size is not None:
                 return self.triton_json_dir+f"/MOE_W4A8INT8[{block_size[0]},{block_size[1]}]_E={E}_N1={N1}_N2={N2}_K={K}_TOPK{TOPK}_{self.device_name}.json"
             else:
                 return self.triton_json_dir+f"/MOE_W4A8INT8_E={E}_N1={N1}_N2={N2}_K={K}_TOPK{TOPK}_{self.device_name}.json"
-        else:
+        elif use_int8_w8a8:
             if block_size is not None:
                 return self.triton_json_dir+f"/MOE_BLOCKINT8[{block_size[0]},{block_size[1]}]_E={E}_N1={N1}_N2={N2}_K={K}_TOPK{TOPK}_{self.device_name}.json"
             else:
                 return self.triton_json_dir+f"/MOE_W8A8INT8_E={E}_N1={N1}_N2={N2}_K={K}_TOPK{TOPK}_{self.device_name}.json"
-
+        else:
+            if block_size is not None:
+                return self.triton_json_dir + f"/MOE_BLOCKFP8[{block_size[0]},{block_size[1]}]_E={E}_N1={N1}_N2={N2}_K={K}_TOPK{TOPK}_{self.device_name}.json"
+            else:
+                return self.triton_json_dir + f"/MOE_W8A8FP8_E={E}_N1={N1}_N2={N2}_K={K}_TOPK{TOPK}_{self.device_name}.json"
 
     def get_moeint8_triton_cache(self,file_path,E,N1,N2,K,TOPK):
+        if file_path in self.cache_json_data:
+            # 直接返回缓存数据，避免重复读取
+            return self.cache_json_data[file_path]
         cache_json_file=file_path
         
         if os.path.exists(file_path):
@@ -2439,7 +2446,7 @@ class W8a8GetCacheJSON:
             for sub_key, sub_value in value.items():
                 configs_key= f"{sub_key}_{key}"   
                 configs_dict[configs_key]=sub_value
-    
+        self.cache_json_data[file_path] = configs_dict
         return configs_dict
    
 # Adapted from: https://stackoverflow.com/a/47212782/5082708