本项目基于 Triton 框架与 Ascend NPU 硬件特性，打造高性能矩阵乘法算子。核心聚焦矩阵乘法的并行效率优化、特殊矩阵适配与参数智能调优，通过创新的 Split-K 双阶段算法，解决了 NPU 片上缓存有限、K 维度并行失效、长瘦矩阵负载不均等关键问题，实现超大矩阵、矩形矩阵及长瘦矩阵场景的高效计算。
项目无需依赖历史版本对比，直接呈现当前最优实现的功能、性能与使用方式，适用于对矩阵乘法性能有高要求的深度学习训练与推理场景。
- Split-K 双阶段并行计算：通过 “并行计算 + 结果归并” 两阶段逻辑，实现 K 维度高效拆分与并行处理，彻底发挥多线程潜力。
- 智能参数自动调优：内置 11 种分块配置组合，结合 Triton Autotune 机制，自动匹配不同矩阵形状与硬件环境的最优参数。
- 长瘦矩阵深度适配：针对 M/N 维度差异显著的矩阵，采用非对称分块策略，优先优化长维度并行效率，抑制短维度调度开销。
- NPU 硬件特性协同：利用 vector 核并行写回、循环展开等技术，掩盖访存延迟，提升内存带宽利用率与硬件流水线效率。
- 全场景矩阵支持：兼容小方阵、大方阵、矩形矩阵、长瘦矩阵等 10 种常见形状，覆盖从实验到生产的各类使用场景。
- 高精度结果保障：与 PyTorch 原生实现结果误差小于 1e-2，满足深度学习计算精度要求。
-
硬件：Ascend NPU（支持 Triton 框架适配）
-
软件：
- Python 3.8+
- Triton 2.0+
- PyTorch 1.18+
- torch_npu（Ascend PyTorch 适配库）
- pandas（性能测试统计）
-
安装 Ascend NPU 驱动及基础软件栈
-
安装依赖包：
pip install triton torch torch_npu pandas -
克隆本项目代码并进入目录：
git clone <项目仓库地址> cd triton-ascend-matmul
import torch
import torch_npu
from matmul_kernel import matmul_wrapper_splitK, KERNEL_V5_SplitK
# 配置 Ascend NPU 设备
torch.npu.set_device(0)
# 生成测试矩阵（超大矩阵/长瘦矩阵推荐使用，效果最佳）
a = torch.randn((4096, 8192), dtype=torch.float16, device="npu")
b = torch.randn((8192, 4096), dtype=torch.float16, device="npu")
# 调用 V5 版 Split-K 优化算子
c = matmul_wrapper_splitK(a, b, kernel_func=KERNEL_V5_SplitK)
# 结果校验（与 PyTorch 原生实现对比）
c_torch = torch.matmul(a, b)
print(f"结果误差: {(c - c_torch).abs().max().item()}") # 误差小于 1e-2
运行内置脚本，自动测试 10 种矩阵形状的性能表现：
python matmul_kernel.py
输出将包含各矩阵形状的平均运行时间、精度校验结果，直观展示算子性能。
| 函数名 | 功能描述 | 参数说明 |
|---|---|---|
matmul_wrapper_splitK(a, b, kernel_func) |
矩阵乘法核心入口 | a/b：输入矩阵（float16 类型，NPU 设备）；kernel_func：指定 V5 优化核（默认 KERNEL_V5_SplitK） |
create_kernel_SplitK(name, autotune_configs, parallel, k_splits) |
自定义 Split-K 算子 | k_splits：K 轴拆分次数（最大支持 4）；parallel：是否启用并行写回 |
run_performance_test() |
性能测试函数 | 自动遍历所有测试矩阵形状，输出性能报告 |
| 矩阵类型 | 具体形状 | 适用场景 |
|---|---|---|
| 小方阵 | 512x512x512 | 轻量级实验、快速验证 |
| 中方阵 | 2048x2048x2048 | 中等规模模型训练 / 推理 |
| 大方阵 | 4096x4096x4096 | 大规模深度学习任务 |
| 矩形矩阵 | 2048x1024x4096、4096x8192x4096 | 特征维度不匹配的计算场景 |
| 长瘦矩阵 | 4096x32x4096、2048x32x2048 | 高维特征提取、稀疏数据计算 |
| 参数名 | 作用 | 最优值 | 设计逻辑 |
|---|---|---|---|
S_UNROLL（循环展开因子） |
提升指令级并行度 | 8 | 平衡 “延迟隐藏” 与 “资源占用”，避免算力浪费或寄存器溢出 |
并行写回分块数 |
优化内存写回效率 | 4 | 适配 Ascend NPU 4 个并行处理单元，最大化内存带宽利用率 |
K_SPLITS（K 轴拆分次数） |
提升 K 维度并行度 | 4 | 基于硬件资源限制设计，平衡并行收益与调度开销 |
AUTOTUNE_CONFIGS（分块配置） |
适配不同矩阵形状 | 11 种组合 | 包含大 / 中 / 小尺寸分块，覆盖全场景矩阵维度需求 |
通过两阶段协同计算，实现 K 维度高效并行，解决传统并行中 “叠加冲突” 问题：
- 阶段一：并行计算：
matmul_splitk_stage1函数按K_SPLITS将 K 维度拆分为多个子任务，每个线程块独立处理部分 K 子块，计算结果写入三维中间缓冲区。 - 阶段二：结果归并：
matmul_splitk_stage2函数遍历所有 K 拆分的部分结果，累加得到完整矩阵，确保计算准确性与并行效率。 - 核心优势：通过中间缓冲区隔离并行任务，避免叠加冲突，让 K 维度并行真正发挥作用。
针对长瘦矩阵（如 M>>N 或 N>>M）的负载不均问题，采用五大核心原则：
- 长轴优先：以 M/N 长轴为核心设计并行策略，最大化长轴缓存利用率。
- 短轴轻量：短轴采用小分块、少并行，减少线程调度与同步开销。
- 非对称分块：分块大小匹配长 / 短轴硬件特性，而非均衡设置。
- 延迟掩盖：通过长轴并行、数据预取、循环展开，解决访存瓶颈。
- 向量复用：通过打包计算调整分块，弥补短轴向量资源浪费。
- 循环展开因子（S_UNROLL=8）：过小则指令并行不足，无法掩盖访存延迟；过大则寄存器资源溢出，导致调度效率下降。
- 并行写回分块数（4）：4 个分块能充分利用 NPU 并行处理单元，同时避免过多分块引发的内存访问冲突。
- 分块配置多样性：11 种分块组合覆盖 32x32x32 到 128x256x64 等尺寸，确保不同矩阵形状都能匹配最优分块。
| 矩阵形状 | 运行时间（ms） | 关键优势 |
|---|---|---|
| 512x512x512 | 0.411 | 精度与原生实现一致，满足轻量计算需求 |
| 2048x2048x2048 | 0.489 | 兼顾速度与资源占用，适合中等规模任务 |
| 4096x4096x4096 | 1.421 | 超大矩阵计算效率突出，充分利用并行资源 |
| 4096x8192x4096 | 1.954 | 矩形矩阵最优场景，K 维度并行优势最大化 |
| 4096x32x4096 | 1.008 | 长瘦矩阵适配有效，避免负载不均导致的性能下降 |
- 超大矩阵（≥4096x4096）：性能最优，并行收益远大于调度开销，充分发挥 Split-K 算法优势。
- 长瘦矩阵（N=32 系列）：相比通用算子，负载均衡性显著提升，避免短维度并行开销浪费。
- 小矩阵（512x512x512）：保证计算精度与功能完整性，适用于快速验证场景。
- 创新 Split-K 双阶段算法，解决 K 维度并行失效问题，并行效率显著提升。
- 全场景适配能力，覆盖 10 种矩阵形状，无需手动调整参数即可高效运行。
- 硬件感知优化，深度结合 Ascend NPU 特性，最大化硬件资源利用率。
- 工程化完善，提供统一 API、性能测试脚本与精度校验逻辑，便于快速集成。
- 小矩阵场景：小矩阵（如 512x512x512）因并行调度开销，性能略低于原生实现。
- 硬件依赖：参数配置针对 Ascend NPU 优化，迁移至其他硬件需重新调优。
- K_SPLITS 上限：受硬件资源限制，K_SPLITS 最大支持 4，无法进一步提升并行度。
- 动态场景适配：新增矩阵形状感知逻辑，自动切换 “并行模式” 与 “轻量模式”，优化小矩阵性能。
- 硬件兼容性扩展：适配更多 NPU/GPU 硬件，提供跨平台参数调优工具。
- 精度扩展：支持 float32、bfloat16 等更多数据类型，满足不同精度需求。
- 功能增强：扩展支持 ReLU、GELU 等激活函数，提升算子通用性。
- 小组成员：王淇辉、袁谦朗、何熙然、娄彦轩
- 项目定位：Ascend NPU 专用高性能矩阵乘法算子，聚焦工业级深度学习计算效率提升
感谢 Triton 框架生态提供的自动调优与并行编程能力，感谢 Ascend NPU 硬件团队的技术支持。本项目基于开源技术与硬件特性优化，欢迎开发者交流反馈与代码贡献！