CANN MXFP4量化矩阵乘教程

MXFP4 量化矩阵乘教程流水分析与分步优化【免费下载链接】cann-samples算子领域高性能实战演进样例与体系化调优知识库项目地址: https://gitcode.com/cann/cann-samples本文档整理各 Step 的问题背景、优化思路与流水对照编译、安装与运行请见matmul_tutorials/README.md。摘要本文档说明matmul_tutorials目录内MXFP4 量化矩阵乘Step 样例的组织方式自Step 0 基准出发经各 Step 优化依次交代各步问题背景、优化思路及与Shape、硬件流水的对照读法。读者可据此选用或排布样例亦可将同一分析范式迁移至其他m×k×n组合。性能建模、Bound 判定与策略取舍的系统性论述见 MX 量化矩阵乘算子性能优化指南。前言各Step技术章节按问题背景 → 优化思路 → 典型 Shape → 示例 Case 与流水分析 → 代码索引比对展开问题背景说明本节承接的实现及其暴露的矛盾即前一阶段在给定假设下仍未能消除的瓶颈优化思路随之给出本节在核内缓冲、多核排布、尾块切分或指令级同步如 UnitFlag等维度上的优化要点。其后各节给出与 Shape 相关的定性条件、可复现m,k,n及 Profiling 对照项。流水与数据流基础本节给出各Step技术章节共用的物理流水视图。存储层次与数据流层级功能Global MemoryGM存放完整 A、B、scaleA、scaleB 及输出 C各个核共用。L1 Buffer位于GM和L0 Buffer之间的一块存储硬件存放A/B矩阵以及Scale矩阵各个核独立。L0 A/B Buffer位于L1 Buffer和MMAD计算单元之间的一块存储硬件存放A/B矩阵各个核独立。L0 MX A/B Buffer位于L1 Buffer和MMAD计算单元之间的一块存储硬件存放Scale A/B矩阵各个核独立。L0C Buffer一块存储硬件存放MMAD计算单元运算的结果各个核独立。数据搬运组件组件功能MTE2GM 与 L1 之间的数据搬运。MTE1L1 与 L0 之间的数据搬运。MMADMXFP4 反量化与矩阵乘累加。FIXPIPEL0C 至 GM 的写回。数据路径可概括为GM →MTE2→ L1 →MTE1→ L0 → MMAD → L0C →FIXPIPE→ GM。符号与约定下列记号通篇统一示例 Case 中的取值须满足 Host 侧校验见各 Step 源码中的参数检查。baseM、baseN核在一个round中完成计算的tile块大小。blockNumDevice 侧的AIC核数。round各个核进行计算的轮数并行计算情况下一轮中32个核会完成32个tile块的计算不满32个tile块则仅有部分核参与计算。优化总览与注意事项下图在四个代表性 Shape 上汇总各优化 Step 的耗时轨迹括号内为相对上一Step的加速变化比例。在完整的代码实现中是否使用A全载模板属于多约束下的联合决策并非由单一静态Shape决定而在tutorial下为简化代码实现与方便演示并未完全实现Tiling侧的逻辑因此图中三条Case无demo7数据。tutorial demo7和recipe A_full_load模板在特定Shape下行为的差异以recipe为准。下方所示Roofline图展现当前性能对比理论性能极限结果。特别说明性能数据相差在1us以内可以认为是硬件执行带来的轻微波动属于正常现象若相邻两步相对性能变化小于阈值则合并为一个桶并标注为demo i - j横坐标AI, FLOPs/Byte $$ AI\frac{2MNK}{Bytes_ABytes_BBytes_{scaleA}Bytes_{scaleB}Bytes_C} $$其中按 MXFP4 教程口径A/B 按0.5 Byte/elemscaleA/scaleB按 K 方向每 32 个元素 1 个 scale1 Byte/elem输出 C 按2 Byte/elem纵坐标Performance, TFLOPS: $$ Perf\frac{2MNK}{duration_us\times 10^{-6}}\times 10^{-12} $$理论Roofline$$ y Min(Bandwidth Peak, Compute Peak) Min(\frac{AI\times BW_{GB/s}}{1000}, Compute Peak) $$0 基础实现0.1 问题背景Step 0 实现一个基础的mxfp4的matmul调用Demo承担基准角色后续的优化策略在此基础上一步一步实施。在尚未引入各优化手段之前按照列优先遍历输出每个Tile块从而完成整个矩阵运算。0.2 实现总结本章节实现mxfp4类型矩阵乘的极简版本 提供标准搬运/MMAD 流水/L1 和L0 单Buffer与通用列优先 block 调度策略给出可复现参考实现Host 启动、列优先Block调度、Block内数据搬运计算时序构成后续各章的对照组。0.3 典型 Shape不涉及0.4 示例 Case 与流水分析Casem8448, k4096, n12288下图是输出一个Tile块的时序图。数据搬运与计算之间存在先后关系。从图中可以看到MMAD的计算依赖MTE1搬运、MTE1的搬运依赖MTE2、FIXPIPE的搬运依赖MMAD的计算。因为本case受L1/L0空间大小限制K方向不可能一次性地载入L0进行计算所以为了简单表示图中只画了2个MMAD指令。0.5 代码索引比对样例根目录0_naive1 Double-Buffer1.1 问题背景从Tile块执行的时序图上进行分析GM→L1 搬运与L1→L0 搬运及 MMAD 计算完全串行流水没有重叠硬件利用率低下。为什么会出现这种情况呢 因为L1、L0A和L0B上只有一块buffer可供使用各条流水线上运行过程中必须要互相等待避免由于数据抢占而导致的精度错误。1.2 优化思路显然L1、L0级上需要开启双Buffer。MTE1读取L1上ping块数据时MTE2写入L1pong块缓存。MTE1读取L1上pong块数据时MTE2写入L1ping块缓存。由此L1 级 Ping-Pong 使 MTE2 与 MTE1 重叠。MMAD计算读取L0A/L0Bping块缓存时MTE1写入L0A/L0Bpong块缓存MMAD计算读取L0A/L0Bpong块缓存时MTE1写入L0A/L0Bping块缓存。由此L0 级 Ping-Pong 使 MTE1 与 MMAD 重叠。未开启PingPong和开启Pingpong的预期流水如下。1.3 典型 Shape不涉及1.4 示例 Case 与流水分析Casem8448, k4096, n12288使用msprof来分析m8448, k4096, n12288的基线性能和ping-pong开启后的性能。从Task时间来看性能显著提高了。 | 场景 | Task时间 | mte2搬运时间 | mte2占比 | cube时间 | cube占比 | |------|---------:|-------------:|---------:|---------:|---------:| | 基线 | 1176.66 | 526.0 | 0.44 | 542.032 | 0.46 | | 开 Ping-Pong 后 | 605.8 | 577.1 | 0.95 | 508.4 | 0.84 |1.5 代码索引比对样例根目录1_pingpong/2 SWAT 调度优化2.1 问题背景承接 Step 0基准为列优先tile 序多核侧若仍采用简单行/列扫描式调度与Step 3本节的 SWAT 调度相比在同一调度轮次内各核所访问的 tile 在M–N 平面上分布范围较大影响首轮搬运与计算阶段的衔接。2.2 优化思路在保持与 Step 0 同类 BlockMmad 数据通路不改 GM→L1→L0→MMAD 流水的前提下仅替换调度侧GetTileIdx引入M 向滑动窗口如WINDOW_LEN4与N 向 Z 型往复使同一轮多核对 tile 的访问尽量集中于较规则的M–N子区域从而提高空间局部性同时使MTE2数据搬运与MMAD的时间重叠程度更高从而提高Cube计算单元利用率。2.3 典型 ShapemCnt×nCnt较大例如二者均不小于 4使得多核同一轮内tile 的空间分布对缓存/带宽行为敏感。k适中或偏大使得单 tile 内核内工作量足以支撑MMAD 连续执行从而放大首轮搬运准备对性能的影响差异。2.4 示例 Case 与流水分析Casem8192, k1024, n1024在baseMbaseN256下totalCntmTileNum×nTileNum32×4128理论上能突出列优先调度与SWAT调度的差异。2.5 代码索引比对样例根目录2_block_swat3 尾轮 tile 负载均衡3.1 问题背景在Step 3中我们对 block 的调度进行了优化。但注意到由于单个 tile 大小在调度阶段确定baseM×baseN及边界tailM/tailN固定的切分策略下尾轮容易出现 tile 块个数小于blockNum核数的情况在此场景下若各 tile 仍作为单块分配至不同核则尾轮可参与的核数会小于blockNum部分核无法参与计算导致尾轮吞吐量下降对于此种情况我们需要考虑一种新的切分策略。3.2 优化思路在沿用 Step 3 式多核排布或同类 SWAT 调度的前提下对尾轮 tile 块实施mTailTile×nTailTile二次切分如图所示减小尾轮 tile 块大小从而增加尾轮 tile 块数量分配到更多的核上计算提高并行度。3.3 典型 Shape尾轮待处理 tile 数显著小于blockNum例如 32 核下尾轮仅有 4 个 tile 块通常保证二次切分后新 tile 数量可增至接近核数。3.4 示例 CaseCasem8448, k4096, n4096在baseMbaseN256下totalCnt33*16528。若blockNum32则round17尾轮余 16 个 tile每个tile大小为256 256如果无尾轮负载均衡则出现16个核计算16个核空闲的情况而在尾轮负载均衡情况下当前固定n方向切分成2块则剩余的16个tile被切分为32个tile每个tile大小为256, 128。从下图中可以看出在对尾轮tile块进行拆分后尾轮中32个核全部得到了利用负载更加均衡最长耗时也相应下降。3.5 代码索引比对样例根目录3_last_round_tile_balance4 UnitFlag 优化4.1 问题背景未开启UnitFlag时FIXPIPE须待MMAD指令全部执行完毕后才能开始将L0C结果搬出。因此MMAD整段计算与FIXPIPE搬出之间是整段完成依赖须先满足全量MMAD结束这一条件FIXPIPE方可启动。在无法用 L0C Double Buffer等机制缓冲、掩盖该依赖时FIXPIPE在MMAD运行期间处于等待两段流水难以时间重叠造成性能损失如图所示。4.2 优化思路UnitFlag为MMAD计算指令与FIXPIPE数据搬运指令提供基于内存访问的细粒度同步512B 粒度。开启后MMAD每计算完 512B 数据FIXPIPE即可搬出对应数据块从而在无法开启 L0C Double-Buffer的条件下提高计算与搬出流水的并行度。开启UnitFlag前后流水对比如下示意图所示4.3 流水分析本节不绑定特定Shape优化前后流水图对比如下所示可以明显看到在UnitFlag优化以后MMAD的执行不在等待FIXPIPE很好的解决了MMAD断流的问题4.4 代码索引比对样例根目录4_unit_flag5 解决 L1 Bank 冲突5_halfl1_ping_halfl1_pong5.1 问题背景本节介绍的优化手段在当前芯片下对整体case性能影响不大。若大家想了解L1的bank冲突相关的技术可以了解下。若不感兴趣只需要了解本优化策略让L1上前一半放入A、scaleA、B和scaleB矩阵的ping块后一半放入A、scaleA、B和scaleB矩阵的pong块。MTE2指令和MTE1指令会同时操作L1的内存前者将pong块/ping块的数据写入L1后者从L1中读取ping块/pong块的数据。芯片策略是优先写L1如果MTE2写入的数据和MTE1读取的数据在同一个bank上那么有可能会受bank冲突影响导致读带宽降低体现在MTE1指令的读取速度上。5.2 优化思路L1上存在2个bank前面一半是和后一半分别是2个bank。如果写数据的地址空间和读数据的地址空间在不同的bank上那么就不会出现bank冲突。当前芯片上bank冲突会影响L1空间上的读取数据。采用顺序 L1 布局A_ping | A_pong | B_ping | B_pong | …时MTE2指令写L1的pong块数据和MTE1指令读取L1的ping块数据时若二者访问的数据落于同一 L1 bank就会产生bank 冲突导致MTE1访问延迟增加。因此采取通过L1前半段放ping块、L1后半段放pong块策略解除 bank 冲突。这样ping块的访问和pong块的访问就能分布在L1两个不同的bank避免bank冲突。下图是 解决L1-BANK冲突前后L1的数据排布图。5.3 典型 ShapeShape[256, 4096, 256]5.4 示例 Case 与流水分析Casem256, k4096, n256拿上面这个 case 来分析解决 bank 冲突前后的 MTE1 指令访问速度。从流水图中统计各 MTE1 指令的平均耗时。MTE1 下主要有两类指令LOAD_2Dv2A、B 的 L1→L0A/L0B与LOAD_MX_2Dv2scaleA、scaleB 的 L1→L0A/L0B。下表为同一 Case 下优化前后对照单位是cycle。场景LOAD_2Dv2 平均耗时LOAD_MX_2Dv2 平均耗时次数解决 bank 冲突前27044232解决 bank 冲突后23338432可见在次数不变均为 32的前提下两类 LOAD 的平均耗时在解决 bank 冲突后均有下降MTE1 搬运效率提高。5.5 代码目录样例根目录5_halfl1_ping_halfl1_pong6 Scale矩阵 内存访问合并优化6_scale_memory_access_coalescing6.1 问题背景现在研究访存bound的case的性能优化。比如[128, 8192, 4096]表现在流水上MTE2 段长时间连续、与后续 MTE1/MMAD 难以充分重叠。对于这种场景优化mte2 的搬运速度才是关键。这类case的简化流水图如下所示本例中一次的搬运量如下可以看出scaleA和scaleB的搬运量小于20K。搬运指令若总数据量小于20K的话对芯片搬运不友好会导致速度降低。$$ \mathrm{baseM}128,\quad \mathrm{baseN}128,\quad \mathrm{baseK}512. $$A左矩阵子块$$ \mathrm{Bytes}{A} \mathrm{baseM}\cdot\mathrm{baseK}\cdot|\mathrm{dtype}{\mathrm{mxfp4}}| 128\times512\times\frac{1}{2} 32768\ \mathrm{B}. $$B右矩阵子块$$ \mathrm{Bytes}{B} \mathrm{baseK}\cdot\mathrm{baseN}\cdot|\mathrm{dtype}{\mathrm{mxfp4}}| 512\times128\times\frac{1}{2} 32768\ \mathrm{B}. $$scaleA沿 K 每 32 个 MXFP4 元素共享 1 个 scale$$ \mathrm{Bytes}{\mathrm{scaleA}} \mathrm{baseM}\cdot\frac{\mathrm{baseK}}{32}\cdot|\mathrm{dtype}{\mathrm{scale}}| 128\times\frac{512}{32}\times 1 2048\ \mathrm{B}. $$scaleB沿 K 每 32 个 MXFP4 元素共享 1 个 scale$$ \mathrm{Bytes}{\mathrm{scaleB}} \mathrm{baseN}\cdot\frac{\mathrm{baseK}}{32}\cdot|\mathrm{dtype}{\mathrm{scale}}| 128\times\frac{512}{32}\times 1 2048\ \mathrm{B}. $$6.2 优化思路为了解决这个问题引入Scale访问合并策略ScaleA 和ScaleB矩阵一次性载入多个 baseK 块代码中使用scaleKL1来表示搬入L1 。在本例的这个case中设置scaleKL1为K方向的整个长度将scale的K方向一次性搬入到L1。这样做的话符合芯片设计规格可以提高访存速度。采用Scale访问合并策略后scaleA和scaleB的搬运量如下。scaleA沿 K 每 32 个 MXFP4 元素共享 1 个 scale$$ \mathrm{Bytes}{\mathrm{scaleA}} \mathrm{baseM}\cdot\frac{\mathrm{scaleKL1}}{32}\cdot|\mathrm{dtype}{\mathrm{scale}}| 128\times\frac{8192}{32}\times 1 32768\ \mathrm{B}. $$scaleB沿 K 每 32 个 MXFP4 元素共享 1 个 scale$$ \mathrm{Bytes}{\mathrm{scaleB}} \mathrm{baseN}\cdot\frac{\mathrm{scaleKL1}}{32}\cdot|\mathrm{dtype}{\mathrm{scale}}| 128\times\frac{8192}{32}\times 1 32768\ \mathrm{B}. $$在本例中采用Scale访问合并策略后scaleA和scaleB只在首次会加载到L1后续只有A和B矩阵加载到L1。6.3 典型 ShapeShape[128, 8192, 4096]6.4 示例 Case 与流水分析我们使用 msprof 工具采集访问合并优化前后的性能。下表为同一 Case 下 使能访问合并优化前后的 Task时间、MTE2时间 与访存占比对照时间单位 μs。通过调整scaleL1的大小来实现配置scaleL1Task 时间MTE2 时间访存占比优化前2×BASE_K102439.5334.40.89优化后16×BASE_K819236.7631.270.87从性能对比来看同样的数据量下 MTE2 搬运时间由 34.4 降低到 31.27 μsTask 总时间相应下降。6.5 代码目录样例根目录6_scale_memory_access_coalescing7 A 矩阵全载优化7_fullload7.1 问题背景上面优化的各个caseA矩阵和B矩阵需要在K方向进行切分一块一块地搬入到L1/L0,完成K方向累加最终将单个Tile块输入到HBM上。这种策略我们叫做非全载策略。但是如果A矩阵比较小而B矩阵特别大若仍然采用上述策略那么肯定会导致A矩阵重复搬运很有可能会导致性能劣化。这回拿[128, 4096, 81920]来举例。使用非全载策略流水示意图如下。矩阵A、scaleA、矩阵B和scaleB矩阵每次载入一个切块到L1/L0,完成K方向累加。让我们计算下左矩阵所占的总空间大小。通过计算发现左矩阵只占L1的一半空间。那么左矩阵是不是可以全载 $$ \mathrm{baseM}128\quad \mathrm{K}4096. $$$$ \mathrm{L1Size} 524288. $$$$ \mathrm{Bytes}{A} \mathrm{baseM}\cdot\mathrm{K}\cdot|\mathrm{dtype}{\mathrm{mxfp4}}| 128\times4096\times\frac{1}{2} 262144\ \mathrm{B}. $$7.2 优化思路因此引入A 矩阵全载A-Full-Load策略将左矩阵 A 常驻在L1空间上。这里采用的方式是按照切块大小去填满A矩阵在L1的空间不是一次性加载完成。scaleA矩阵是一次性加载到L1的。右矩阵B和scaleB矩阵是按照切块大小一块一块地载入到L1。每计算输出一个切块都可以复用驻留在L1上数据这样可以最大限度地减少数据的重复搬运。采用该策略的流水示意图如下。7.3 典型 ShapeShape[128, 4096, 81920]7.4 示例 Case 与流水分析我们使用 msprof 工具采集非全载策略与全载策略下 Shape[128, 4096, 81920]的性能数据。下表在scaleKL1KL11024相同切片设定下对比两种策略时间单位 μs。策略Task 时间MTE2 时间访存占比非全载183.87178.350.97全载170.07164.250.97全载策略下 Task 与 MTE2 耗时均低于非全载访存占比数值相同。下面估算两种策略下MTE2 等效搬运带宽字节/周期。时间与周期换算芯片上1 μs 1650 cycle。单核在全局上负责的B与scaleB搬运量整块问题、与 A 重复次数无关为$$ \mathrm{baseM}128,\quad \mathrm{baseN}512,\quad \mathrm{baseK}512,\quad K4096. $$$$ \mathrm{blockNum}32,\quad N81920. $$$$ \begin{aligned} \mathrm{Bytes}{B}^{(\mathrm{pc})} \frac{N,K,|\mathrm{dtype}{\mathrm{mxfp4}}|}{\mathrm{blockNum}} \ \frac{81920\times 4096\times \frac{1}{2}}{32} \ 5242880\ \mathrm{B}. \end{aligned} $$$$ \begin{aligned} \mathrm{Bytes}{\mathrm{scaleB}}^{(\mathrm{pc})} \frac{N}{\mathrm{blockNum}}\cdot\frac{K}{32}\cdot|\mathrm{dtype}{\mathrm{scale}}| \ \frac{81920}{32}\times\frac{4096}{32}\times 1 \ 327680\ \mathrm{B}. \end{aligned} $$非全载时左矩阵沿N按切块推进单核在N方向需经历的块次数即A / scaleA 在全 (K) 上需重复搬运的次数为$$ \begin{aligned} T \frac{N}{\mathrm{blockNum}\cdot\mathrm{baseN}} \ \frac{81920}{32\times 512} \ 5. \end{aligned} $$非全载A、scaleA 在每个N切块搬运时需要陪同搬运一次故$$ \begin{aligned} \mathrm{Bytes}{A}^{\mathrm{nl}} T\cdot \mathrm{baseM},K,|\mathrm{dtype}{\mathrm{mxfp4}}| \ 5\times 128\times 4096\times \frac{1}{2} \ 1310720\ \mathrm{B}. \end{aligned} $$$$ \begin{aligned} \mathrm{Bytes}{\mathrm{scaleA}}^{\mathrm{nl}} T\cdot \mathrm{baseM},K\cdot\frac{|\mathrm{dtype}{\mathrm{scale}}|}{32} \ 5\times 128\times 4096\times \frac{1}{32} \ 81920\ \mathrm{B}. \end{aligned} $$MTE2 总搬运量与等效速度为$$ \begin{aligned} \mathrm{Bytes}{\mathrm{total}}^{\mathrm{nl}} \mathrm{Bytes}{B}^{(\mathrm{pc})}\mathrm{Bytes}_{\mathrm{scaleB}}^{(\mathrm{pc})}\mathrm{Bytes}_{A}^{\mathrm{nl}}\mathrm{Bytes}_{\mathrm{scaleA}}^{\mathrm{nl}} \ 5242880 327680 1310720 81920 \ 6963200\ \mathrm{B}. \end{aligned} $$$$ \begin{aligned} v_{\mathrm{MTE2}}^{\mathrm{nl}} \frac{\mathrm{Bytes}{\mathrm{total}}^{\mathrm{nl}}}{T{\mathrm{MTE2}}\times 1650} \ \frac{6963200}{178.35\times 1650} \ \approx 23.66\ \mathrm{B/cycle}. \end{aligned} $$全载A、scaleA 各驻留 L1仅搬运一次故$$ \begin{aligned} \mathrm{Bytes}{A}^{\mathrm{fl}} \mathrm{baseM},K,|\mathrm{dtype}{\mathrm{mxfp4}}| \ 128\times 4096\times \frac{1}{2} \ 262144\ \mathrm{B}. \end{aligned} $$$$ \begin{aligned} \mathrm{Bytes}{\mathrm{scaleA}}^{\mathrm{fl}} \mathrm{baseM},K\cdot\frac{|\mathrm{dtype}{\mathrm{scale}}|}{32} \ 128\times 4096\times \frac{1}{32} \ 16384\ \mathrm{B}. \end{aligned} $$$$ \begin{aligned} \mathrm{Bytes}{\mathrm{total}}^{\mathrm{fl}} \mathrm{Bytes}{B}^{(\mathrm{pc})}\mathrm{Bytes}_{\mathrm{scaleB}}^{(\mathrm{pc})}\mathrm{Bytes}_{A}^{\mathrm{fl}}\mathrm{Bytes}_{\mathrm{scaleA}}^{\mathrm{fl}} \ 5242880 327680 262144 16384 \ 5849088\ \mathrm{B}. \end{aligned} $$$$ \begin{aligned} v_{\mathrm{MTE2}}^{\mathrm{fl}} \frac{\mathrm{Bytes}{\mathrm{total}}^{\mathrm{fl}}}{T{\mathrm{MTE2}}\times 1650} \ \frac{5849088}{164.25\times 1650} \ \approx 21.58\ \mathrm{B/cycle}. \end{aligned} $$非全载与全载的搬运速度同一量级全载通过去掉 A、scaleA 的重复搬运在总搬运量更小的前提下得到更短 MTE2 时间与更优端到端性能。二者对比说明在访存 bound 场景下减少左矩阵重复搬运能有效提升实际表现。7.5 代码目录样例根目录7_fullload【免费下载链接】cann-samples算子领域高性能实战演进样例与体系化调优知识库项目地址: https://gitcode.com/cann/cann-samples创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考