import os import sys import time import json import subprocess import platform sys.path.insert(0, os.path.join(os.path.dirname(__file__), "..", "build")) from kt_kernel import kt_kernel_ext import torch from tqdm import tqdm # 测试参数设置 expert_num = 256 hidden_size = 7168 intermediate_size = 2048 m_block = 1 group_min_len = 10 group_max_len = 1024 num_experts_per_tok = 8 # layer_num = 5 # 测试时不同的层数 # qlen = 1 # warm_up_iter = 100 # test_iter = 10000 layer_num = 1 # 测试时不同的层数 qlen = 1024 warm_up_iter = 100 test_iter = 10000 CPUINFER_PARAM = 304 # 初始化 CPUInfer（此处使用原始构造函数，可根据需要调整配置参数） CPUInfer = kt_kernel_ext.CPUInfer(CPUINFER_PARAM) # 获取脚本相关信息，用于生成结果保存文件名 script_path = os.path.abspath(__file__) script_dir = os.path.dirname(script_path) script_name = os.path.splitext(os.path.basename(script_path))[0] json_path = os.path.join(script_dir, "bench_results " + ".jsonl") def get_git_commit(): """ 获取当前 git 提交记录（commit hash 和提交信息）， 并检查是否存在未提交的更改（dirty） """ result = {} try: commit = subprocess.check_output(["git", "rev-parse", "HEAD"]).decode("utf-8").strip() commit_msg = subprocess.check_output(["git", "log", "-1", "--pretty=%B"]).decode("utf-8").strip() result["commit"] = commit result["commit_message"] = commit_msg # 检查是否存在未提交的更改 dirty_output = subprocess.check_output(["git", "status", "--porcelain"]).decode("utf-8").strip() if dirty_output: result["dirty"] = True result["dirty_files"] = dirty_output.splitlines() else: result["dirty"] = False except Exception as e: result["commit"] = None result["commit_message"] = None result["dirty"] = None result["error"] = str(e) return result def get_system_info(): """ 获取系统信息，包括系统名称、CPU 型号、内存大小（GB）、CPU 核数及 socket 数量 """ info = {} uname = platform.uname() info["system_name"] = uname.system info["node_name"] = uname.node # 获取 CPU 型号（仅 Linux 支持） cpu_model = None if os.path.exists("/proc/cpuinfo"): try: with open("/proc/cpuinfo", "r") as f: for line in f: if "model name" in line: cpu_model = line.split(":", 1)[1].strip() break except Exception as e: cpu_model = f"Error: {e}" info["cpu_model"] = cpu_model # 获取内存大小（单位：GB），仅 Linux 支持 mem_total_gb = None if os.path.exists("/proc/meminfo"): try: with open("/proc/meminfo", "r") as f: for line in f: if "MemTotal" in line: mem_kb = float(line.split(":", 1)[1].split()[0]) mem_total_gb = round(mem_kb / (1024 * 1024), 2) break except Exception as e: mem_total_gb = f"Error: {e}" info["memory_size_GB"] = mem_total_gb info["cpu_core_count"] = os.cpu_count() # 解析 /proc/cpuinfo 获取 socket 数量 sockets = set() if os.path.exists("/proc/cpuinfo"): try: with open("/proc/cpuinfo", "r") as f: for line in f: if "physical id" in line: sockets.add(line.split(":", 1)[1].strip()) except Exception as e: sockets = set() info["cpu_socket_count"] = len(sockets) if len(sockets) > 0 else 1 return info def record_results(result, filename=json_path): """ 将结果以 JSON 格式追加到文件中 """ with open(filename, "a") as f: f.write(json.dumps(result) + "\n") def bench_moe(quant_mode: str): """ 依据不同量化模式进行 MoE 性能测试，包含预热与测试阶段 """ with torch.inference_mode(): # 根据量化模式设置数据类型与 bytes_per_elem hidden_type = 30 # ggml_type::GGML_TYPE_BF16（固定） if quant_mode == "fp32": gate_type = 0 # ggml_type::GGML_TYPE_F32 up_type = 0 down_type = 0 bytes_per_elem = 4.0 elif quant_mode == "fp16": gate_type = 1 # ggml_type::GGML_TYPE_F16 up_type = 1 down_type = 1 bytes_per_elem = 2.0 elif quant_mode == "bf16": gate_type = 30 # ggml_type::GGML_TYPE_BF16 up_type = 30 down_type = 30 bytes_per_elem = 2.0 elif quant_mode == "q8_0": gate_type = 8 # ggml_type::GGML_TYPE_Q8_0 up_type = 8 down_type = 8 bytes_per_elem = 1.062500 elif quant_mode == "q6_k": gate_type = 14 # ggml_type::GGML_TYPE_Q6_K up_type = 14 down_type = 14 bytes_per_elem = 0.820312 elif quant_mode == "q5_k_m": gate_type = 13 # ggml_type::GGML_TYPE_Q5_K up_type = 13 down_type = 14 # ggml_type::GGML_TYPE_Q6_K bytes_per_elem = 0.731771 elif quant_mode == "q4_k_m": gate_type = 12 # ggml_type::GGML_TYPE_Q4_K up_type = 12 down_type = 14 # ggml_type::GGML_TYPE_Q6_K bytes_per_elem = 0.648437 elif quant_mode == "q3_k_m": gate_type = 11 # ggml_type::GGML_TYPE_Q3_K up_type = 11 down_type = 13 # ggml_type::GGML_TYPE_Q5_K bytes_per_elem = 0.515625 elif quant_mode == "q2_k": gate_type = 10 # ggml_type::GGML_TYPE_Q2_K up_type = 10 down_type = 11 # ggml_type::GGML_TYPE_Q3_K bytes_per_elem = 0.328125 elif quant_mode == "iq3_xs": gate_type = 21 # ggml_type::GGML_TYPE_IQ3_S up_type = 21 down_type = 21 bytes_per_elem = 0.429688 elif quant_mode == "iq2_xxs": gate_type = 16 # ggml_type::GGML_TYPE_IQ2_XXS up_type = 16 down_type = 16 bytes_per_elem = 0.257812 else: raise ValueError("不支持的量化模式") # 构建各层 MoE 模型 moes = [] for _ in tqdm(range(layer_num), desc="Initializing MOEs"): gate_proj = ( torch.randn((expert_num, intermediate_size, hidden_size), dtype=torch.float16, device="cpu") .to("cpu") .contiguous() ) up_proj = ( torch.randn((expert_num, intermediate_size, hidden_size), dtype=torch.float16, device="cpu") .to("cpu") .contiguous() ) down_proj = ( torch.randn((expert_num, hidden_size, intermediate_size), dtype=torch.float16, device="cpu") .to("cpu") .contiguous() ) config = kt_kernel_ext.moe.MOEConfig(expert_num, num_experts_per_tok, hidden_size, intermediate_size) config.pool = CPUInfer.backend_ config.m_block = m_block config.group_min_len = group_min_len config.group_max_len = group_max_len config.gate_proj = gate_proj.data_ptr() config.up_proj = up_proj.data_ptr() config.down_proj = down_proj.data_ptr() config.gate_type = gate_type config.up_type = up_type config.down_type = down_type config.hidden_type = hidden_type moe = kt_kernel_ext.moe.MOE(config) CPUInfer.submit(moe.load_weights_task()) CPUInfer.sync() moes.append(moe) # 生成输入数据 print("Generating data...") # 专家路由索引与权重，每层一个 gen_iter = 1000 expert_ids = ( torch.rand(gen_iter * qlen, expert_num, device="cpu") .argsort(dim=-1)[:, :num_experts_per_tok] .reshape(gen_iter, qlen * num_experts_per_tok) .contiguous() ) weights = torch.rand((gen_iter, qlen, num_experts_per_tok), dtype=torch.float32, device="cpu").contiguous() input_tensor = torch.randn((layer_num, qlen, hidden_size), dtype=torch.bfloat16, device="cpu").contiguous() output_tensor = torch.empty((layer_num, qlen, hidden_size), dtype=torch.bfloat16, device="cpu").contiguous() # 将 qlen 封装成 tensor，用于 forward 调用 qlen_tensor = torch.tensor([qlen], dtype=torch.int32) # 预热阶段 print("Warming up...") for i in tqdm(range(warm_up_iter), desc="Warm-up"): CPUInfer.submit( moes[i % layer_num].forward_task( qlen_tensor.data_ptr(), num_experts_per_tok, expert_ids[i % gen_iter].data_ptr(), weights[i % gen_iter].data_ptr(), input_tensor[i % layer_num].data_ptr(), output_tensor[i % layer_num].data_ptr(), False, ) ) CPUInfer.sync() # 测试阶段 print("Start testing...") start = time.perf_counter() for i in tqdm(range(test_iter), desc="Testing"): CPUInfer.submit( moes[i % layer_num].forward_task( qlen_tensor.data_ptr(), num_experts_per_tok, expert_ids[i % gen_iter].data_ptr(), weights[i % gen_iter].data_ptr(), input_tensor[i % layer_num].data_ptr(), output_tensor[i % layer_num].data_ptr(), False, ) ) CPUInfer.sync() end = time.perf_counter() total_time = end - start # 计算性能指标 time_per_iter_us = total_time / test_iter * 1e6 bandwidth = ( hidden_size * intermediate_size * 3 * num_experts_per_tok * (1 / 8 * 256 * (1 - (31 / 32) ** qlen)) * bytes_per_elem * test_iter / total_time / 1e9 ) # 单位：GB/s flops = ( hidden_size * intermediate_size * qlen * 3 * num_experts_per_tok * 2 * test_iter / total_time / 1e12 ) # 单位：TFLOPS # 打印结果 print("Quant mode:", quant_mode) print("Time(s):", total_time) print("Iteration:", test_iter) print("Time(us) per iteration:", time_per_iter_us) print("Bandwidth:", bandwidth, "GB/s") print("TFLOPS:", flops) print("") # 整理测试结果 result = { "test_name": os.path.basename(__file__), "quant_mode": quant_mode, "total_time_seconds": total_time, "iterations": test_iter, "time_per_iteration_us": time_per_iter_us, "bandwidth_GBs": bandwidth, "flops_TFLOPS": flops, "timestamp": time.strftime("%Y-%m-%d %H:%M:%S", time.localtime()), "test_parameters": { "expert_num": expert_num, "hidden_size": hidden_size, "intermediate_size": intermediate_size, "m_block": m_block, "group_min_len": group_min_len, "group_max_len": group_max_len, "num_experts_per_tok": num_experts_per_tok, "layer_num": layer_num, "qlen": qlen, "warm_up_iter": warm_up_iter, "test_iter": test_iter, "CPUInfer_parameter": CPUINFER_PARAM, }, } # 添加 git 与系统信息 result.update(get_git_commit()) result.update(get_system_info()) # 将结果记录到 JSON 文件中 record_results(result) if __name__ == "__main__": # 根据需要选择量化模式，目前调用 q4_k_m 模式，对 layer_nums 列表中各层数进行测试 bench_moe("q4_k_m") # 其他量化模式调用可以按需取消注释 # bench_moe("fp32", layer_num) # bench_moe("fp16", layer_num) # bench_moe("bf16", layer_num) # bench_moe("q8_0") # bench_moe("q6_k", layer_num) # bench_moe("q5_k_m", layer_num) # bench_moe("q3_k_m", layer_num) # bench_moe("q2_k", layer_num) # bench_moe("iq3_xs", layer_num) # bench_moe("iq2_xxs", layer_num)