当前，新型AI芯片架构加速崛起，英伟达CUDA生态“高墙”正被逐步瓦解。作为国产新型架构算力芯片代表企业，清微智能携手智源研究院，积极参与国产AI生态建设。
近日，清微技术团队依托FlagTree的TLE扩展，成功落地Compute-Shift GEMM计算模式。
本文将详细展示清微智能如何基于Triton-TLE，实现Compute-Shift GEMM计算模式，通过将硬件特性自然融入高层编程模型，让开发者用熟悉的语言，轻松释放RPU芯片的极致算力，最终实现2.5倍性能提升！
FlagTree统一AI编译器：FlagOS的核心组件，构建起支持Triton语言、面向多种AI硬件架构的增强型编译器，通过统一代码仓库与多后端机制，实现“一次开发，跨芯迁移”。
在此基础上推出的Triton-TLE（Triton Language Extensions）语言扩展体系，在保持与标准Triton兼容的同时，提供三个层次的语言扩展，将单点语言能力拓展至兼顾高性能、高开发效率及多架构适配的广阔空间。
01.
瓶颈：
当传统GEMM算子
遇上可重构数据流架构
清微的RPU（可重构处理单元）采用了一种独特的循环存储与计算相结合的数据流架构。其核心组件包括：
RT（Reconfigurable Tile）： 图中的 T 表示一个 RT。每个 RT 内部包含用于 Tensor 和 Vector 计算的 CGRA 单元，以及一部分 Scratchpad Memory（SPM）。RT 是 RPU 的基本计算与存储单元，不同 RT 之间可以进行数据通信。
NoC（Network on Chip）: NoC 负责 RT、高速 IO 和存储设备之间的片上互联，是数据传输的核心通道。
LPDDR : LPDDR 是片外存储器，用于在片上存储空间不足时提供更大的容量支持。
High-speed IO : High-speed IO 用于实现芯片之间的高速互联，支持跨芯片的数据传输与协同计算。
RPU的核心是RT，每个RT包含独立的计算单元（CGRA）和SPM，RT之间通过NoC高速互联，形成分布式存储 + 分布式计算的体系结构。
这种分布式架构虽然优势明显，但也对算子实现提出了新的要求。由于每个RT的SPM容量有限，对于大 shape的GEMM往往需要多次计算，这就导致需要反复从LPDDR搬运数据，使得片外带宽成为性能瓶颈。
考虑一种最朴素的情况，一个AxB->C的矩阵乘， 如果只对左矩阵进行空间维度上的切分，则大概情况如下图所示：每个RT各自加载一整份右矩阵到本地SPM。若有16个RT，右矩阵对LPDDR的读取量就被放大了16倍。
LPDDR 带宽迅速成为瓶颈，片上算力再强也难以充分释放。问题的本质在于：数据复用被忽视，RT间的高速互联未被用于削减外部访存。
02.
破局：
Compute-Shift GEMM的
计算与通信协同
要突破这一瓶颈，就需要根据架构特点调整计算模式——让GEMM算子以RT协同的方式展开，每个RT在本地完成一部分计算的同时，将下一步所需的数据分片传递给其他RT。
这种将本地计算与RT间数据流转结合起来的执行方式，我们称之为Compute-Shift GEMM。其核心思想是 “compute + shift”协同进行，即一边计算、一边通信，从而提升片上数据复用率并降低外部访存压力。具体流程如下：
Step 0: 每个Tile Load一份右矩阵的子块， 计算得到当前输出的一部分；
Step 1: 将当前加载的右矩阵子块通过NoC发送到下一个RT；
Step 2: 接收到上一个Tile发送到数据，和已加载到左矩阵做Compute得到部分输出；
Step 3: 重复上述行为，直到一个RT获取完所有的右矩阵子块，完整流程如下：
这种模式用片上通信换取了外部访存的大幅削减，尤其适合大Shape场景，能够有效缓解LPDDR带宽压力。
03.
落地：
Triton-TLE如何将
硬件特性融入
高层编程模型
实现Compute-Shift GEMM需要解决一个核心问题：如何在高层编程模型中自然地描述“RT 间数据流动”，而不陷入底层硬件细节？而这正是Triton-TLE的价值所在。
Triton-TLE 三层接口
Triton-TLE 提供了三种架构的接口抽象，满足不同背景的开发者需求：
TLE-Lite：对 Triton 的轻量级扩展，所有特性兼容各类硬件后端，仅需对原有 Triton kernels 少量修改即可拿到大幅性能提升。主要面向算法工程师和快速性能优化场景。
TLE-Struct：按硬件的架构聚类抽象，面向不同分类（如 GPGPU、DSA）提供扩展，满足进一步性能优化的需求。需要开发人员对目标硬件的特性和优化技巧有一定了解。
TLE-Raw：提供对硬件最直接的控制，支持使用硬件厂商的原生编程语言获取最极致的性能。需要开发人员对目标硬件的深入了解，主要面向性能优化专家。
其中 TLE-Lite 和 TLE-Struct 会通过 FLIR最终 Lowering到 LLVM IR，而TLE-Raw 则通过语言对应的编译管线(如厂商的私有编译器)Lowering到 LLVM IR。最后它们会被 Link到一起，共同生成一个完整的 kernel供 Runtime 加载和执行。
核心原语：从“标记”到“通信”
Triton-TLE目前已实现了部分公用接口用于做通信模式的描述， 为了实现上述的Compute-Shift GEMM计算模式， 清微智能复用了Triton-TLE 中已有的通信接口语义，并对其后端Lowering做了合适的适配，主要采用到的接口大致有：
• tle.dsa.alloc(...)
• tle.dsa.local_ptr(...)
• tle.remote( )...
• tl.store(...) 搭配remotepointertle.dsa.alloc
• tle.dsa.alloc
tle.dsa.alloc用于在spm上分配一块buffer， 后续用于存储接收到的数据。
tle.remote
用于声明一块buffer为远端buffer（所有tile都分配了相同的一块buffer），返回相应的写地址
延迟实例化： 远程信息只是作为元数据附着在 Python对象上，真正的IR生成推迟到后续调用 local_ptr() 时才发生。
tle.remote() 不生成任何 IR op。它是一个纯 Python 层的标记操作：
1) 对 buffered_tensor 做浅拷贝；
2) 在拷贝的 .type 上设置
_tle_remote_shard_id = send_next_tile；
3) 在拷贝的 .type 上设置
_tle_remote_scope = mesh；
4) 返回标记后的 buffered_tensor。
tle.dsa.local_ptr
根据标记选择不同的create分支
前面提过，remote会标记tensor属性（remote_buffer_marker），若local_ptr的输入tensor：
• 无remote_buffer_marker
获取本地alloc出来的buffer的地址
• 有remote_buffer_marker
获取远端读写地址，通过tl.store可以直接触发通信
编译流程：从高层语义到硬件指令的降级
编译流程：从高层语义到硬件指令的降级
Triton-TLE 的编译器后端完成了关键的“降级”过程，这是整个方案能够“用统一语言驾驭异构硬件”的技术基石：
• Remote只对tensor进行标记，后续local_ptr识别到被标记的tensor会额外生成remote_pointers去修饰local_pointers；
• remote_pointers+tl.store的组合会被合并优化为mk.remote_store。
通过这种设计，开发者用“声明式”的代码描述了RT间的数据流转，而编译器负责将其精确映射到底层硬件。远程通信不再需要显式调用发送原语，而是通过“对远程 buffer 的 store 操作”自然触发——这正是Triton-TLE抽象能力的关键体现。
Kernel 实现：循环移位结构的代码表达
完整的Compute-Shift GEMM kernel 采用循环移位结构，代码逻辑清晰体现了“计算-发送-接收-移位”的流水：
04.
收益：
性能实测与数据解读
通过在清微RPU上对Compute-Shift GEMM进行实测，并与标准Triton实现及清微原生算子进行对比，得出以下结论：
• 相比标准Triton原生实现，Triton-TLE版本性能提升达2.5倍；
• 相比清微原生Compute-Shift手工优化版本，达到其性能的1.12倍，证明Triton-TLE的表达能力已接近手调极致性能。
尽管通信会带来额外开销，但随着数据量增大，Compute-Shift在减少访存上的收益愈发显著，是大规模深度学习任务的理想方案。
清微RPU所采用的可重构数据流架构，正是面向大算力场景的重要硬件方向。
而Triton-TLE的出现，让行业不再需要为每一种硬件架构维护一套独立的底层代码库。通过Compute-Shift GEMM这一典型案例，切实看到：
• Triton-TLE Struct层接口能够以简洁的语义表达RT间通信与计算并行；
• “标记 + 延迟实例化”的编译流程设计，让远程通信自然融入store操作，开发者无需手动管理NoC传输细节；
• 编译器后端的合并优化，将高层语义精确降级为硬件指令，在保持代码可读性的同时逼近手调性能上限。
05.
实力：
稳居国产AI生态
第一梯队
早在去年6月，清微智能已正式成为智源研究院战略合作伙伴单位，深度参与 FlagOS 开源生态建设。
目前，在众多芯片适配厂商中，清微智能的适配模块数量位列第四，在非GPGPU架构芯片中排名第二，充分彰显了新型架构在开源生态中的兼容性与成长潜力，成为国内AI生态第一梯队选手。
清微技术团队已完成适配的FlagOS核心模块包括：
这一系列适配成果，标志着清微智能与智源研究院的合作已从“单点适配”走向“全栈融合”。双方正以开源共建的方式，为国产AI生态探索一条“硬件+软件”深度协同的可行路径。
当前，中国AI行业正从“破局”向“立势”跨越。未来，清微智能将继续携手行业伙伴，持续为国产AI生态贡献坚实创新力量。
1
2
热门跟贴