TPU：TPU SparseCore

可视化演示

演示TPU SparseCore的gather/scatter架构、嵌入表稀疏查找优化及稀疏vs稠密计算性能对比，交互切换模式与稀疏度。

稀疏计算困境

小智走进的第一间库房，就是存放“用户偏好表”的地方。这座库房有上百万个格子，每个格子里放着一份向量档案——这就是嵌入表（Embedding Table），记录着每个用户、商品、视频的特征。理论上，AI 模型在推理时要查这张表，把用户和物品变成可以相乘计算的向量。

问题在于访问模式。一次推荐请求，往往只需要查几十到几百个 ID 的档案，就像读者进馆只借三五本书。可是传统的矩阵计算单元（如 TPU 上的 MXU、GPU 上的 Tensor Core）是为“稠密矩阵乘法”设计的——它们要求把整个表的所有行都读出来，排成规则的大方块，才肯开始算。于是每次推理，硬件都要把上百万行从 HBM 显存浩浩荡荡搬一遍，绝大多数行根本没用到就被搬进计算单元，再原封不动搬回去。

这场“满表搬家”造成的访存浪费，在每次只查少量 ID 的推荐、搜索场景中触目惊心。真正完成有用的计算只占搬运量极小比例，访存带宽成了整个系统的咽喉。以一次查询 1000 个 ID、表共 1000 万行为例：

小智明白了：这座百万格子的库房，每次只为借几本书就把所有书搬一遍，自然慢得离谱。要提速，必须换一种思路——不搬整个书架，让硬件学会“按索引取书”。这正是 SparseCore 登场的理由。

SparseCore架构

图书馆请来了一位经验丰富的分拣主管，他把整套流程拆成了三个角色来分工。取书员（访存引擎）拿着读者递来的借书单（ID 列表），飞奔到库房里只把对得上号的书取出来；借阅台（嵌入片）负责把书中的关键段落（向量）做计算处理；最后，汇总处（汇聚片）把多路取回的资料合并成完整的答复。

在 SparseCore 中，这三个角色是硬件电路层面真实存在的模块。访存引擎（Access Engine）的核心任务是高效索引——它根据输入的稀疏 ID，直接计算出各行在存储中的物理地址并发出访问请求，不再遍历全表。嵌入片（Embedding Tile）是计算主力的所在地，它把从存储器取回的向量行与权重做乘加运算，类似于借阅台在有限资料内完成深度分析。汇聚片则把从多个嵌入片返回的碎片化结果按用户 ID 聚合起来，还原成完整的输出张量。

三个角色按硬件流水线（Pipeline）方式协同工作。访存引擎去取“下一批”数据时，嵌入片同时在算“这一批”，汇聚片则在整理“上一批”的结果——像图书馆的取书、翻书、登记三条流水线并行运转，谁也不等谁。流水线让访存延迟被计算时间掩盖。

分拣主管的这套流水线让小智大受启发：只要让取书员、借阅台、登记处并行干活，库房再大也不怕。但小智还注意到一件事：取书员能在百万格子里瞬间定位，光靠跑得快还不够——仓库里的书必须按照某种巧妙的方式摆放。这就引出了压缩存储格式。

压缩存储格式

为了取书员能飞速定位，分拣主管重整了库房：不再把百万本书整整齐齐占满整个楼层，而是把所有书压缩进一排紧凑的书架上，并单独做一本“索引册”，记录每本书在书架上的实际位置。这本索引册，对应到 SparseCore 中就是 CSR（压缩稀疏行）格式的核心思想。

CSR 格式用两个数组描述一张稀疏表：索引数组（Indices）按顺序记录每本“有内容的书”的原始 ID；数值数组（Values）紧挨着存放与索引一一对应的向量。当读者报出一个 ID，取书员只需要在索引册里用二分查找或哈希定位到该 ID 的位置，再顺着位置在同一“书架层”取出相邻数据即可。由于稀疏行不再按原始 ID 大排列布，而是压缩成连续存储，局部性大大提高——一次取一个 ID 往往能顺带把附近几个 ID 的数据一起搬进缓存。

更妙的是“就地更新”（In-place Update）机制。推荐系统每来一批新点击，都要更新对应行的向量。在传统满表计算中，更新一个行往往要先把整个表拷到新地址再整体写回；而在 SparseCore 中，更新只发生在 CSR 紧凑区域内，直接找到行号覆盖写入，不需要移动其他数据。压缩后的存储开销显著下降：

小智看着这本索引册感慨：百万本书被压成几排书架，索引册一翻就能找到，存储省了 99% 都不稀奇。不过他又想到一层——哪怕取书员再快，如果每次都要跑到仓库深处去取，速度终究有限。要是能在借阅台旁边放一个“热门书柜”就好了，这就轮到准入缓存登场。

准入缓存机制

分拣主管显然早有准备。他在借阅台旁增设了一个“热门书柜”（准入缓存），专门存放最近高频借阅的书。读者来借书时，先看热门书柜里有没有，有就直接递过去；没有才让取书员跑仓库。这个热门书柜不需要太大，放最常访问的几百本，就能挡住绝大部分取书跑腿。

硬件上，这个热门书柜就是 SparseCore 中的准入缓存（Access Cache）。它位于核心与主存之间，容量远比主存小但访问速度极快，专门缓存近期被访问过的嵌入行。其核心指标是命中率 $h$——所有访问请求中，能在缓存里直接满足的比例。若访问模式呈“二八分布”（20% 的热门行承接 80% 的请求），那么只要缓存能容纳下这部分热门行，命中率就能冲到 80% 甚至更高。

但缓存不是无脑吸收所有数据。SparseCore 设计了一个“准入队列”（Admission Queue），只有被判定为“值得缓存”的行才被放入缓存，防止那些只被访问一次就永不出现的冷门行把热门行挤出去。更关键的设计是“幂等性”（Idempotency）：同一份访存请求即使被重复执行多次，得到的结果也完全一致，不会因为缓存中数据状态新旧而产生混乱——就像热门书柜里的一本书无论被借多少次，归还时内容都不会被改坏。这保证了缓存替换策略可以大胆工作而不必担心数据一致性崩溃。命中率提升带来的总延时变化：

小智把热门书柜的故事记在笔记本上：缓存就是“让 80% 的借书需求在几步之内解决”。不过他很快又想到一个更深的问题——仓库里的书在不断变化：新书补充、旧书被淘汰、书被借走又还回，长期下来书架上会留下大量空洞（碎片）。谁来收拾这个烂摊子？分拣主管微微一笑，指了指调度室里的仓库管理系统。

内存管理策略

仓库管理系统（内存管理）是分拣主管的另一项秘密武器。它负责两件事：一是给新入库的书分配书位（单元分配），二是把那些已经无人问津、长期占据书位的书进行回收（隐式回收）。在 SparseCore 中，存储器被划分为固定大小的“单元”（Cells），每个单元可容纳一个嵌入向量。分配器维护着一张空闲单元表，取书员需要空间时以 O(1) 复杂度从中取一个空闲单元。

隐式回收则更像图书馆的“还书检测系统”：当某个嵌入行被覆盖更新时，旧的向量值如果没有其他引用，就会被自动标记为可回收。SparseCore 不依赖软件垃圾回收器，而是用硬件计数器追踪每个单元的最后使用时间，超时未用的单元会被硬件悄悄回收进空闲库。这套机制让内存管理的开销几乎被隐藏在正常访存流程中。

长期运行后，仓库难免出现碎片——一些单元被释放后与相邻空闲不连续，形同书架上的空洞。此时系统会计算碎片率 $f$，若碎片率超过阈值，则触发一次“紧凑整理”，把碎片中间的存活单元移动到连续区域，腾出大片完整空间。整理操作是昂贵的，所以控制在碎片率不足以影响大块分配时才触发，实现“懒整理”。

小智终于理解，SparseCore 之所以敢把主存让给稀疏表随便用，正是因为背后有一套自动化的“书位管家”。现在他迫不及待想知道：这套取书、缓存、整理体系，究竟让真实世界的哪些应用焕然一新？

工作负载加速

小智打开 TPU 的测试报告，发现受益最大的恰好就是他做的那类系统。推荐系统是非常典型的工作负载：嵌入表有数亿行，每次推理只查几十个槽位（slot）对应的行。在 SparseCore 上，推荐模型不用再被迫把整个嵌入表压实进计算单元，而是“查一行算一行”，配合准入缓存的超高命中率，吞吐量比传统稠密执行方式提升数倍。

图神经网络（GNN）是第二类大赢家。GNN 做邻居聚合时，每个节点要访问其邻居节点的特征行，而真实世界的图往往是幂律分布的——少数“明星节点”拥有海量邻居，大多数节点的邻居寥寥无几。SparseCore 对这类不规则访存模式尤其擅长：访存引擎为每个节点并行取邻居行，嵌入片就地聚合，汇聚片再把不同节点子图的结果拼接。过去 GPU 上因邻居数参差不齐导致的线程发散和访存浪费，在这里被硬件流水线吸收。

除了这两种场景，Transformer 中的稀疏注意力、大规模 MoE（专家混合）模型中专家路由等也呈现出稀疏特性。SparseCore 的设计哲学是“多类稀疏共优化”——只要有“大表 + 少数索引”的访问规律，就能压缩存储、缓存提速、流水线隐藏延迟。综合性能提升可以用加速比衡量：

合上测试报告，小智长长舒了一口气。当年那个让他在深夜抓狂的推荐系统，如今在 SparseCore 上跑得飞快。他这才彻底明白：面对“书极多、借得极少”的世界，聪明的方法从来不是像传统硬件那样把图书馆整个搬空，而是派一个懂索引的取书员，备一个放热门书的书柜，再雇一个细心的书位管家——只为需要的行搬书，就是 SparseCore 的全部秘密。

AI 应用场景

Python 演示

import random class DenseEmbedding:    """传统稠密嵌入表：预分配全部行，模拟满表计算"""    def __init__(self, vocab_size: int, dim: int):        self.vocab_size = vocab_size        self.rows = [[0.0] * dim for _ in range(vocab_size)]     def lookup(self, ids):        return [self.rows[i][:] for i in ids]  class SparseEmbedding:    """稀疏嵌入表：按需分配行，模拟 SparseCore 的惰性分配"""    def __init__(self, vocab_size: int, dim: int):        self.dim = dim        self.rows = {}     def lookup(self, ids):        result = []        for i in ids:            if i not in self.rows:                self.rows[i] = [0.0] * self.dim            result.append(self.rows[i])        return result  def simulate():    vocab_size = 100000    dim = 32    batch_ids = [random.randint(0, vocab_size - 1) for _ in range(500)]     dense = DenseEmbedding(vocab_size, dim)    sparse = SparseEmbedding(vocab_size, dim)     dense_rows = dense.lookup(batch_ids)    sparse_rows = sparse.lookup(batch_ids)     dense_size = vocab_size * dim    sparse_size = sum(len(v) for v in sparse.rows.values())    print('稠密表元素数:', dense_size)    print('稀疏表实际元素数:', sparse_size)    print('内存节省比例: {:.2%}'.format(1 - sparse_size / dense_size))    print('访问行数/总行数: {}/{}'.format(len(sparse.rows), vocab_size))  if __name__ == '__main__':    simulate()

1. self.rows = [[0.0] * dim for _ in range(vocab_size)]稠密表预先把 10 万行全部建好，模拟传统 GPU 满表常驻显存的做法。
2. self.rows = {}稀疏表用字典只保存真正被访问过的行，模拟 SparseCore 的按需分配。
3. if i not in self.rows: self.rows[i] = [0.0] * self.dim首次访问某行时才分配内存，这就是隐式分配的精髓。
4. batch_ids = [random.randint(0, vocab_size - 1) for _ in range(500)]生成 500 个随机查询 ID，模拟一批推荐请求的稀疏访问。
5. sparse_size = sum(len(v) for v in sparse.rows.values())统计稀疏表实际占用的元素总数，通常远小于稠密表。
6. print('内存节省比例: {:.2%}'.format(...))直观输出相对稠密表节省的显存百分比，通常在 95% 以上。

常见误区

面试题