引言在 AI 大模型时代算子性能优化是提升整体训练和推理效率的关键。TileLang是一门面向高性能算子开发的领域特定语言DSL采用简洁直观的编程范式让开发者能够以接近数学表达的方式描述计算逻辑。相比传统的手写算子开发TileLang 大幅降低了开发门槛使开发者能够更高效地完成高性能算子的开发与调优。TileLang-Ascend是 TileLang 针对 Ascend NPU 的高性能算子开发框架在保持易用性的同时提供了强大的性能优化能力。通过内置的优化原语如T.pipelined、T.tile等和自动化优化策略开发者可以轻松实现核间流水、核内并行、内存常驻等高级优化技术无需手动编写复杂的同步和调度代码。这种高层抽象 底层优化的设计理念使得在 TileLang-Ascend 上进行性能优化既高效又便捷。本文将系统介绍 TileLang-Ascend 框架下的算子性能优化方法论并以 Flash Attention (FA) 和 Sparse Flash Attention (SFA) 算子为例展示优化思路的实际应用与显著效果。一、性能优化方法1.1 先把单核做到足够快再谈多核协同在 TileLang-Ascend 里算子性能优化天然分两层C/V 核内优化单核是不是够快和C/V 核间优化多核能不能高效并行。为什么一般选择先做核内、再做核间原因很实际如果单核本身就有大量气泡和等待核间流水线再怎么设计都会被拖住——慢流水拖快流水优化空间还没打开就先被锁死了。反过来单核打满之后核间流水的效果才能完整释放。优化层次核心问题典型手段C/V核内Intra-Core单核的计算和搬运能不能充分重叠Cube 核 L1 常驻、多 Buffer、Vector 核num_stages调优、scalar 向量化C/V核间Inter-CoreCube 和 Vector 之间流水能不能接紧num_stages调优、任务均衡、同步优化1.2 优化不是一个动作而是一个闭环性能优化不是选几个手段叠上去而是测→找→改→验的迭代闭环性能基准测试 → 识别瓶颈 → 针对性优化 → 验证收益 ↑ └────────── 迭代 ←──────────┘具体四步建立基线通过 profiling 获取初始性能数据明确当前版本的耗时分布。定位瓶颈判断瓶颈主要来自核内计算、数据搬运还是核间同步与调度。实施优化围绕主要瓶颈选择最合适的手段避免一次性叠加过多优化影响问题定位。验证收益对比优化前后的性能变化确认收益是否稳定同时检查是否引入新的瓶颈。迭代几轮之后性能会收敛到一个比较稳的状态。1.3 三条原则贯穿始终在具体执行过程中建议始终围绕以下三个原则展开原则说明掩盖短流水将耗时短的流水尽量用耗时长的流水掩盖减少气泡优化任务排布减少核间等待时间收敛到单一 Bound理想情况下优化至单一类型流水 bound2. 核内优化2.1 多Buffer搬运和计算别再串着跑Cube或Vector都可以开启多buffer以vector为例计算队列、MTE2搬运队列、MTE3搬运队列互相独立天然可以并行。多 Buffer 就是利用这个特性把数据分成多份让 CopyIn、Compute、CopyOut 三个阶段交错执行以double buffer为例串行模式: ——————————时间轴—————————— [MTE2][VEC][MTE3][MTE2][VEC][MTE3] Double Buffer: ——————————时间轴—————————— buffer0: [MTE2][VEC][MTE3] buffer1: [MTE2][VEC][MTE3]原理很简单buffer0和buffer1使用不同的地址空间当 buffer0 做 Compute 的时候buffer1 就可以进行 CopyIn 搬入下一块数据搬运时间被藏在计算后面Vector 单元的等待时间直接降下来。也可以根据循环次数已经数据总量选择更多的buffer数但随着buffer数增加占用的空间也会增加。2.2 Cube 核L1 常驻省空间流水争取优化到单一BoundCube 核做矩阵乘的时候L1 容量有限Q/K/V矩阵不一定全放得下。这时候可以选择让复用率最高的那份长驻 L1只分批搬运其他矩阵减少 GM↔L1 之间的搬运次数。以 FA 算子为例常见策略策略适用条件做法大复用L1 充裕Q 在 L1 中持续多个基本块每次只换 KV小复用L1 紧张Q 只保留一个基本块省空间具体选哪份常驻要看算子的访问模式也可以反过来让 KV 常驻效果类似。常驻做完之后下一步就是收敛到单一 Bound。理想状态下Cube 核上应该只剩一条流水是瓶颈耗时最长的那条其余流水全部被它盖住优化前各流水串行: 时间轴 → MTE: [] [] M: [] [] FIX: [] 优化后FIX 作为 BoundMTE 和 M 被它掩盖: 时间轴 → |------ 一个FIX周期 ------| MTE: [] ← 被 FIX 盖住 M: [] ← 被 FIX 盖住 FIX: [] ← 瓶颈流水当 MTE搬运和 M矩阵乘都能被 FIX 阶段盖住的时候系统就只剩一个瓶颈针对瓶颈流水当前为FIX优化才能有进一步的性能提升这个判断非常关键它决定了优化往哪个方向推才是有效的。2.3 Vector 核向量化 指令合并Vector 核经常遇到的瓶颈是 scalar 操作过多Scalar 负责标量运算和流程控制它会阻塞整个异步并行流水所以优化方向很明确把 scalar 循环改成 tile 操作让一批数据一次处理完。# 优化前对 m_i 逐元素 scalar 操作forh_iinrange(block_M//2):T.tile.sub(acc_s_ub[h_i,:],acc_s_ub[h_i,:],m_i[h_i])# 优化后一次 broadcast 一次 tile 操作T.tile.broadcast(m_i_2d,m_i,tmp_ub)T.tile.sub(acc_s_ub,acc_s_ub,m_i_2d)scalar 减完之后还可以进一步合并可融合的指令减少指令下发次数# 优化前两条指令T.tile.mul(acc_s_ub,acc_s_ub,sm_scale)T.tile.sub(acc_s_ub,acc_s_ub,m_i_2d)# 优化后一条 axpy 搞定T.tile.axpy(acc_s_ub,m_i_2d,sm_scale)指令下发次数直接影响 Vector 单元的调度开销。尤其在高频循环里少一条指令、每轮就少一个调度节拍累积下来对性能提升的帮助也很大。3. 核间优化提升 Cube 与 Vector 的流水线衔接效率核内做到位之后下一个瓶颈通常会移到 Cube 和 Vector 之间的衔接上。CV 融合算子的核间优化建议按以下顺序推进先把慢的核内部流水做好 → 选合适的流水深度 → 降低核间同步开销。3.1 num_stages 调优控制流水线深度以减少气泡T.Pipelined是 TileLang-Ascend 的流水线循环原语用来把普通循环组织成可重叠执行的producer/consumer流水。在 C/V 核间流水场景下Cube 核通常作为producer将中间结果写入workspaceVector 核作为consumer从workspace读取结果并继续处理。num_stages表示producer和consumer之间可使用的最大buffer数决定这条流水线可同时保留多少个中间结果版本。合适的num_stages能提高前后任务的重叠度减少Cube与Vector之间的等待时间。对于CV融合算子当参与流水的任务块数较多且C核与V核执行时间不均衡时适当增大num_stages可以减少气泡3.2 核间同步优化核间同步过于频繁会增加Scalar控制开销容易出现scalar bound但同步次数过少又可能降低C/V核间并行度带来性能损失。应在保证正确性的前提下权衡同步开销与并行收益选择较优的同步次数# 优化前每次任务都同步foriinrange(n):process()sync()# 优化后多次任务后同步foriinrange(n):process()ifi%21:sync()如果确认核间同步的耗时较大可以增大cross_interval参数查看性能收益起始频率一般可以从每 2 次任务同步一次开始试然后根据 profiling 调整原则和前面的迭代闭环一样改一个变量测一轮数据确认方向对了再继续。四、实操案例一Flash Attention 算子优化Flash Attention 是 Transformer 模型中的核心算子其性能直接影响模型训练和推理效率。4.1 性能测试输入参数定义参数取值说明B1Batch大小Q_N12Query序列长度KV_N1KV序列数D128隐藏层维度S32K/64K/128K序列长度block_size128块大小最佳性能结果SAscend CTileLang性能百分比Ascend C/TileLang32K37555u46643u80.52%64K149578u185188u80.77%128K600018u741211u80.95%4.2 优化策略及收益分析针对 FA 算子我们采用了以下优化组合L1 内存常驻Q 矩阵在 L1 中持续多个基本块减少 GM 访问指令向量化将 scalar 操作转换为 tile 操作减少scalar指令— 原生Ascend C算子36%多Buffer核内流水并行掩盖数据搬运延迟核内冗余同步消除— 原生Ascend C算子50%T.pipelined 原语开启核间 CV 流水最大化 Cube 和 Vector 核并行度调整num_stages为8— 原生Ascend C算子60%优化核间同步下发次数每两次任务进行一次核间同步— 原生Ascend C算子72%减少指令数使用axpy代替mul和sub减少指令下发数— 原生Ascend C算子80%4.3 优化效果通过系统性优化FA 算子在保持 TileLang 的高开发效率前提下达到了 Ascend C 原生算子80%的性能混合编程模式下性能达到原生算子60%。优化项L1 内存常驻指令向量化多Buffer核内冗余同步消除CV pipelined优化核间同步下发次数减少指令数性能A3flash_attn_bhsd_expert_h16_d128.py√√√√√√√80%flash_attn_bhsd_auto_pipeline_h16_d128.py√√√×√×√60%算子实现https://github.com/tile-ai/tilelang-ascend/tree/ascendc_pto/examples/flash_attention/fa_opt五、实操案例二Sparse Flash Attention 算子优化Sparse Flash Attention (SFA) 是 DeepSeek v3.2 版本中引入的核心注意力机制。本节展示从基线版本到极致优化版本的完整优化路径。5.1 性能测试输入参数定义T1, B1, Q_N128, KV_N1D512, D_rope64sparse_size2048, block_size128三组测试 shape仅在 KV_S 上不同编号KV_S说明shape02560短序列shape16400中序列shape248000长序列5.2 优化策略及收益分析以kv_s2560的shape为例第一阶段baseline → developer性能提升 42%通过 Broadcast 优化将多条指令合并减少指令下发开销。同时针对离散访存将kv gather为连续内存后搬出性能从 602us 提升至 347us。第二阶段developer → T.pipelined性能提升 63%引入稀疏访存优化和 CV pipeline利用T.pipelined原语开启核间流水实现 Cube 和 Vector 核的高效并行。性能从 347us 提升至 127us。第三阶段T.pipelined → no_cv_pipeline性能提升 14%通过增大 s2 切分大小64→256增加每轮 Gather kv 和 Copy out 的数据规模有利于提升访存吞吐同时增大基本块大小也能进一步提升cube侧计算效率。。性能从 127us 提升至 109us。5.3 优化效果从基线版本到最优版本总体性能提升达 82%602us → 109us充分验证了系统化优化方法论的有效性。在tileLang-ascend仓上我们提供了四个逐步优化的版本文件Fixed Cores2切分大小kv gather与连续搬出异步拷贝/核内手动同步搬入搬出 Ping-PongCV pipelineBroadcast优化AXPY优化性能数据A3sparse_flash_attn_pa_baseline.py√64××××××602ussparse_flash_attn_pa_developer.py√64√×××√×347ussparse_flash_attn_pa.py√64√√√√√√127ussparse_flash_attn_pa_no_cv_pipeline.py√256√√√×√√109us算子实现https://github.com/tile-ai/tilelang-ascend/tree/ascendc_pto/examples/sparse_flash_attention/bench_sfa六、性能调优 Checklist在实际优化过程中建议按以下清单逐项检查采集 msprof 性能数据分析 C/V 核耗时比例尝试不同num_stages值检查 L1/L0 内存利用率确认 Double Buffer 已开启优化 scalar 操作为向量化减少不必要的核间同步七、常见问题与解决方案现象可能原因解决方案C 核大量气泡V 核耗时长num_stages 太小增大 num_stages内存溢出num_stages 过大或 buffer 过大减小分块参数指令下发慢scalar 操作过多使用 T.tile 向量化GM 带宽未打满数据搬运效率低开启 L1 常驻、Double Buffer结语算子性能优化是一项系统性工程需要从核内和核间两个维度协同发力。本文通过 FA 和 SFA 两个实际案例展示了从方法论到实践的完整优化路径。TileLang-Ascend在帮助开发者提高开发效率的同时也提供了必备的性能调优能力。欢迎更多开发者加入高性能算子的建设。TileLang 社区已开放完整示例、性能脚本与优化指南欢迎体验与贡献。相关资源TileLang-Ascend Programming GuideTileLang-Ascend关键特性文档MindStudio Insight 下载