EngineX-Cambricon/enginex-mlu370-vllm
11
0
Fork 1
Code Issues Pull Requests Actions Projects Releases Wiki Activity

add ops

Browse Source
This commit is contained in:
Chranos
2026-02-04 17:39:32 +08:00
parent 8511fe8530
commit 79dfc69789
299 changed files with 55927 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files

89
torch_mlu_ops-v1.3.2/benchmarks/benchmark_ffn.py Normal file
View File

@@ -0,0 +1,89 @@
import torch
import torch_mlu
import torch_mlu_ops as tmo
from common import benchmark_forward, save_to_csv
import argparse
from tabulate import tabulate
import os
e2e_time_param_dict_list = [{"batch": 1, "seq_len": 1024, "hidden_size": 1600, "inner_size": 6400,
"gated_ffn": False, "act_mode": "gelu", "input_dtype": [torch.float16, torch.bfloat16]},
{"batch": 1, "seq_len": 1024, "hidden_size": 2048, "inner_size": 8192,
"gated_ffn": True, "act_mode": "gelu", "input_dtype": [torch.float16, torch.bfloat16]},
{"batch": 1, "seq_len": 1024, "hidden_size": 4096, "inner_size": 11008,
"gated_ffn": True, "act_mode": "gelu", "input_dtype": [torch.float16, torch.bfloat16]},
{"batch": 1, "seq_len": 1024, "hidden_size": 4096, "inner_size": 14336,
"gated_ffn": True, "act_mode": "gelu", "input_dtype": [torch.float16, torch.bfloat16]},
{"batch": 1, "seq_len": 1024, "hidden_size": 4096, "inner_size": 16384,
"gated_ffn": True, "act_mode": "gelu", "input_dtype": [torch.float16, torch.bfloat16]},
{"batch": 1, "seq_len": 1024, "hidden_size": 5120, "inner_size": 13824,
"gated_ffn": True, "act_mode": "gelu", "input_dtype": [torch.float16, torch.bfloat16]},
{"batch": 1, "seq_len": 1024, "hidden_size": 5120, "inner_size": 27392,
"gated_ffn": True, "act_mode": "gelu", "input_dtype": [torch.float16, torch.bfloat16]},
{"batch": 1, "seq_len": 1024, "hidden_size": 6656, "inner_size": 17920,
"gated_ffn": True, "act_mode": "gelu", "input_dtype": [torch.float16, torch.bfloat16]},
{"batch": 1, "seq_len": 1024, "hidden_size": 8192, "inner_size": 22016,
"gated_ffn": True, "act_mode": "gelu", "input_dtype": [torch.float16, torch.bfloat16]},
{"batch": 1, "seq_len": 1024, "hidden_size": 8192, "inner_size": 24576,
"gated_ffn": True, "act_mode": "gelu", "input_dtype": [torch.float16, torch.bfloat16]},
{"batch": 1, "seq_len": 1024, "hidden_size": 8192, "inner_size": 28672,
"gated_ffn": True, "act_mode": "gelu", "input_dtype": [torch.float16, torch.bfloat16]},
{"batch": 1, "seq_len": 1024, "hidden_size": 8192, "inner_size": 49152,
"gated_ffn": True, "act_mode": "gelu", "input_dtype": [torch.float16, torch.bfloat16]},
{"batch": 1, "seq_len": 1024, "hidden_size": 12288, "inner_size": 32768,
"gated_ffn": True, "act_mode": "gelu", "input_dtype": [torch.float16, torch.bfloat16]},
{"batch": 1, "seq_len": 1024, "hidden_size": 14336, "inner_size": 57344,
"gated_ffn": True, "act_mode": "gelu", "input_dtype": [torch.float16, torch.bfloat16]}]
def main():
parser = argparse.ArgumentParser()
parser.add_argument('--repeat_times', type=int, default=10, help='repeat times for testing')
parser.add_argument('--csv', action='store_true', help='write the report data to csv')
parser.add_argument('-o', type=str, help='specify the output folder name under --csv mode')
args = parser.parse_args()
device = 'mlu'
titles = ["batch", "seq_len", "hidden_size", "inner_size", "gated_ffn", "act_mode", "input_dtype", "hardware_time(us)", "e2e_latency(us)"]
contents = []
for params_dict in e2e_time_param_dict_list:
batch = params_dict["batch"]
seq_len = params_dict["seq_len"]
hidden_size = params_dict["hidden_size"]
inner_size = params_dict["inner_size"]
gated_ffn = params_dict["gated_ffn"]
act_mode = params_dict["act_mode"]
input_dtype_list = params_dict["input_dtype"]
for dtype in input_dtype_list:
if dtype == torch.bfloat16 and not torch_mlu.mlu.is_bf16_supported():
continue
input = torch.randn(batch, seq_len, hidden_size).to(device).to(dtype)
up_proj_weight = torch.randn(inner_size, hidden_size).to(device).to(dtype)
up_proj_bias = torch.randn(inner_size).to(device).to(dtype)
down_proj_weight = torch.randn(hidden_size, inner_size).to(device).to(dtype)
down_proj_bias = torch.randn(hidden_size).to(device).to(dtype)
gate_up_proj_weight, gate_up_proj_bias = None, None
if gated_ffn:
gate_up_proj_weight = torch.randn(inner_size, hidden_size).to(device).to(dtype)
gate_up_proj_bias = torch.randn(inner_size).to(device).to(dtype)
hardware_time, e2e_time = benchmark_forward(tmo.ffn,
input,
up_proj_weight,
up_proj_bias,
down_proj_weight,
down_proj_bias,
gate_up_proj_weight,
gate_up_proj_bias,
act_mode,
repeats=args.repeat_times)
content = [f"{batch}", f"{seq_len}", f"{hidden_size}", f"{inner_size}", f"{gated_ffn}", f"{act_mode}", f"{dtype}", f"{hardware_time}", f"{e2e_time}"]
contents.append(content)
table = [titles] + contents
print(tabulate(table, headers="firstrow", tablefmt="grid"))
if args.csv:
current_file_path = __file__
_, file_name = os.path.split(current_file_path)
save_to_csv(table, args.o, file_name)
if __name__=="__main__":
main()