前阵子我手头有个模型上线项目batch里的Embedding查询耗时一直压不下去profile一看居然能占到整个推理的18%。我翻了一阵子调用栈最后把矛头指向那个看起来人畜无害的哈希操作——把稀疏的ID映射成稠密的Embedding行号。这个动作在CPU上只是调一个unordered_map但放到某AI加速卡的NN算子库环境里就成了一个必须从头设计的高性能哈希算子。这篇文章就聊聊我在这个算子库里实现哈希算子时做过的取舍、踩过的坑和最终跑出来的效果希望对准备碰硬件算子或者做稀疏模型加速的朋友有点参考价值。1. 为什么一个看似简单的哈希值得单独写算子1.1 先从一次Embedding查找的耗时说起CTR预估、推荐召回、图模型中到处都有Embedding层。最典型的流程是线上拿到一批稀疏特征ID每个ID去Embedding表里查对应的稠密向量然后拼接、池化、进神经网络。Embedding表有多大呢几百万行、每行128维float很常见也就是一张表几百MB甚至上GB。在传统实现里ID到行号的映射通常放在CPU侧。原因很简单Python里有dictC里有unordered_map现成的哈希表随手就能用。但你只要在真实推理链路里测一次就会发现这条链路有三个隐藏开销。第一个开销是CPU哈希查找本身。一批请求进来可能是几万个ID每个ID要做一次字符串或整数到行号的映射。虽然单个查找就几纳秒到几十纳秒但几万个堆在一起、还有cache miss累计下来就是毫秒级。第二个开销是数据搬运。CPU查完ID得到的是行号行号只是一个int数组但Embedding向量本体在加速卡侧。你要先把这个行号数组从主机搬到卡上卡再去读Embedding表。这个H2D拷贝虽然不大在低延迟场景里依然扎眼。第三个开销更隐蔽Embedding表如果按行号排列那行号分布是随机的卡侧读取Embedding时会大量踩中cache miss。CPU哈希查找已经让ID访问变成随机访问而卡侧取Embedding又随机一次双重随机访问让访存系统很难受。我当时就想能不能把ID到行号的哈希表直接搬进加速卡算子内部让整条链路不出设备这就需要一个能在算子里高效运行、支持批量查询和批量插入的哈希结构。传统库函数根本没法原样塞进去所以哈希算子必须自己写。1.2 通用哈希为什么帮不上忙不少人第一反应是哈希表不就是一个数组加一个哈希函数吗直接开个全局表在算子代码里逐个查找不就完了。真上手会发现通用哈希库和硬件算子之间的鸿沟很大。先看依赖。std::unordered_map依赖标准库、依赖动态内存分配、依赖全局运行时状态。加速卡算子的编译链路通常要尽量精简很多运行时功能根本不可用。你不太可能在算子内部new一个节点出来再挂在某个桶的链表上那不仅慢还可能直接编译不过。再看并发模型。CPU哈希表是单机多线程模型一个线程处理一批key锁或原子操作保护表结构。但在加速卡上成百上千个核同时涌进来每个核都拿着自己的ID列表哈希表是全局共享的如果还是用锁那性能直接崩到没法看。更麻烦的是有些加速卡虽然支持原子操作但语义和x86稍有差异比如有些只支持整型原子不保证浮点原子有些对未对齐访问会直接报错。最后是内存访问模式。GPU或NPU上性能好坏极度依赖访存局部性。通用哈希表的节点散落在内存各处每个key访问一串指针天然就是访存灾难。如果你希望把哈希表压进片上缓存或高带宽内存数据布局必须自己设计这已经超出了“用现成库”的范畴。所以结论很简单如果想在推理路径上把哈希这一步也塞进端到端设备侧执行就必须写一个定制哈希算子。它不是那种“随便实现一下就行”的工具函数而是需要针对硬件特点重新设计的数据结构。2. 哈希函数、桶数量与冲突策略设计层面的关键取舍2.1 哈希函数不是越随机越好市面上哈希函数一大把MD5、SHA系列、CityHash、MurmurHash、FNV但真正能在算子核函数里跑的其实没那么多。密码学哈希比如SHA系列计算开销太大一个key进去要做几十轮位运算拿来做Embedding查找完全是浪费。我最后用的是类似MurmurHash3 finalizer和lowbias32混合的思路。原理并不复杂输入是一个32位整数ID把它做几次移位、异或、乘法让结果分布尽量均匀。这个计算在加速卡上很友好乘法和移位是基础指令一个ID大概只需要几条指令比密码学哈希快一个数量级以上。给个参考实现大家在别的平台上也很容易验证uint32_t hash_id(uint32_t x) { x ^ x 16; x * 0x7feb352dU; x ^ x 15; x * 0x846ca68bU; x ^ x 16; return x; }把结果跟桶数量掩码做与运算就能得到桶下标前提是桶数量取2的幂。这里有一个取舍要注意哈希函数不是越随机越好。密码学意义上的强随机性通常意味着高计算复杂度。在算子场景里我们只需要满足两个要求一是分布足够均匀不出现大量key撞同一个桶二是确定性同一个ID在任何一次运行中必须算出同样的桶号。后者对调试和离线复现特别重要因为模型上线后的行为应该是可预期的。2.2 桶数量与内存布局桶数量我一开始觉得越大越好毕竟桶多了冲突少。但实测下来并非如此这个我后面会单独说。合理做法是先估算活跃ID数量再把桶数设成2的幂让负载因子落在0.7到0.8之间。假设活跃ID约100万那桶数就取2^21也就是2097152。这样每个桶的负载不到0.5理论上冲突很少。但实际选多少最终要结合内存访问模式来定。桶数太大会让表在内存里铺得太开每个桶大概率只被访问一次缓存命中率反而下降。内存布局上我强烈建议把key数组和value数组分开按结构体数组SoA而不是数组结构体AoS来排布。原因很简单查找时第一步要比较key而value在命中之后才用。如果每个slot里是“keyvalue”紧挨着那查一次key就强迫加载整条valuevalue如果是一个行号还好但如果后面放Embedding向量那一个缓存行里就装不下几个slot了。单独放key可以让一次缓存行加载覆盖更多桶。下面的对比能看出差别布局方式查找时加载内容缓存利用率写入复杂度AoSkey紧挨valuekey value一起加载命中前浪费带宽写起来简单SoAkey和value分开只加载key数组同样一级缓存可查更多桶命中后再返回value实际测试里把value换成行号后SoA比AoS快了大概20%到30%。当value是Embedding向量时这个差距会大到不可忽视。2.3 冲突策略链式还是开放寻址哈希表绕不开冲突处理。常见的三种方案我都比过拉链法、线性探测、Robin Hood哈希。拉链法在CPU库里很常见每个桶维护一个链表。但到了算子内部就不太现实因为链表节点需要动态分配而加速卡端通常没有好用的内存分配器。就算预先分配节点池多核并发往节点池里申请内存也会引入新的竞争点。线性探测实现最简单完全基于数组如果桶被占了就沿着下一个槽位找。它的优点是局部性好因为探测范围都落在附近几个连续slot里非常符合缓存行加载逻辑。缺点是容易聚类一旦某个区域填得比较满新的key会跨越很长一串。Robin Hood是线性探测的改进版它会动态调整每个key的位置让查找次数尽量均衡。写入时如果遇到比自己离理想桶更远的key就把对方挤走自己占住这个位置。实现比线性探测复杂但在高负载因子下能明显减少最坏情况查找长度。我的结论是如果负载因子控制在0.8以下线性探测已经足够如果业务上允许高负载因子Robin Hood更稳。为了并发插入方便最终选了线性探测。原因在于Robin Hood的“挤走”操作在并发场景下需要更复杂的原子协议而线性探测配合原子CAS就能满足大部分需求。3. 从单核到多核并行化设计与流水线落地3.1 多核分片与原子操作算子性能的瓶颈通常不在哈希计算而在并发写入。想象一下几百个核同时拿到一批ID它们各自算出哈希结果之后可能好几个核都命中同一个空位想同时插入。如果不用原子操作后写覆盖先写那就会丢key。我用的方案是把输入ID数组按核数切成连续区间每个核只处理自己那段。哈希表和输出数组是全局共享的但每个核访问位置由哈希结果决定。插入时使用原子CAS保证同一个slot只有一个写入者成功if (atomic_cas(keys[slot], EMPTY, key) EMPTY) { values[slot] row_id; out[i] row_id; } else { // 冲突继续探测 }这里必须非常注意顺序先写value再通过CAS发布key。因为其他核看到某个slot从EMPTY变成非EMPTY后就会认为这个slot已经可用了。如果先CAS把key写进去value还没来得及写其他核读到的就是未初始化的value这个bug非常难追。有些加速卡的原子CAS只支持32位整数对指针或者64位ID无能为力。如果业务ID本身就是64位比较麻烦的一点是你没法用单个32位原子去判断slot是否为空。我当时的做法是把key和value都限制在32位Embedding行号本来也用不了那么高ID则预先做一次压缩或者直接用自增后的int32 ID。这样整个表只操作32位原子简单可靠。3.2 访存优化与数据搬移哈希算子如果每次访问都等主存那性能一定不行。加速卡通常有专门的高带宽内存但带宽依然是最稀缺的资源。我能做的优化主要有三个。第一key数组单独连续排布。前面已经说过连续内存让CPU或加速卡能用缓存行预取一下。哈希查找是随机访问但桶内探测的连续几个slot还是可能落在同一个缓存行里这样查一个key时顺便把相邻几个slot的状态都加载进来一次访存覆盖多次探测。第二把表尽量放到访问延迟更低的内存域。不是所有加速卡都允许你在算子核函数里指定存储类型但如果有片上SRAM或者一级缓存可以放部分热点slot效果会非常明显。我做过一个实验把哈希表最热的1%区域单独缓存一份结果平均查找耗时下降了大概12%。不过这要看算子运行时是否支持这种手动管理不能硬套。第三批量输出合并回写。每个key查到一个value后不要马上分散写回全局输出数组而是在核内暂存一段连续结果攒够一批之后一次性做批量写。这样做能显著减少写放大尤其当输出值很小、ID顺序随机的时候。3.3 一条完整的算子执行流写出来的执行流水线大致如下主机端把原始ID数组搬运到设备侧按核数均分并给每个核发一个start/end区间。每个核先把属于自己的那一段ID加载到高速暂存区避免反复访问全局输入数组。核心循环逐个计算哈希映射到桶线性探测查找或插入。命中时把value写到输出数组的对应位置未命中时分配一个新的行号并插入表。全部核处理完后发一个同步信号再统一把输出行号数组交给后续算子。这里有一个比较反直觉的细节插入新行号时row_id的分配也是个共享计数器多核同时分配就会冲突。我采用的方案是每个核先本地预留一段连续的行号区间比如每个核预留4096个用完再原子fetch_add。这样共享原子操作的次数从每个key一次降到每4096个key一次性能提升明显。3.4 实测数据下面这组数据来自我手头环境下的一次对比测试配置细节不展开了重点看相对趋势。输入是6.4万个ID活跃ID约100万桶数2^21负载因子约0.5实现方式单次batch平均耗时CPU unordered_map H2D拷贝0.96ms单核朴素哈希算子0.82ms四核并行 原子CAS0.27ms四核 SoA布局 批量写入0.16ms最终版含行号预分配0.14ms可以看到只把hash挪到设备侧还不够单核版本对比CPU优势不大。真正的收益来自并行化、内存布局优化和批量写入。最终版本比原先的CPUH2D方案快了接近7倍。4. 调优过程中踩过的坑与排查链路4.1 坑一原子操作语义差异导致计数丢失最开始我做的是一个带频次统计的哈希算子每次命中key后value里的计数要加1。简单写就是uint32_t old values[slot]; values[slot] old 1;这个代码在单核下没问题四核并发时计数会时不时少掉几个。排查过程很典型我先是print每个核内部的局部计数发现每个核都认为自己的写入成功了但最终总数对不上。后来怀疑是values被不同核同时读改写于是改成原子加。问题并没有完全解决。换成原子加之后数值还是会偶发不一致只是频率低了。进一步查文档和实验才发现这个平台上的原子加只保证操作本身原子不保证后续读到的一定是最新值如果你不做release语义的fence其他核可能在缓存里看到旧值。解决办法是在每个核的批量处理结束后加一个内存屏障或者对输出结果做一次“所有核同步后重读”。这个坑给我们的教训是原子操作在硬件上不是魔法语义细节和x86并不完全一样出问题时要怀疑docs里没写的部分而不是默认“原子一定可靠”。4.2 坑二开放寻址在低负载因子下反而慢我把表容量设成活跃ID的四倍觉得冲突会很少结果性能反而下降了。当时很奇怪冲突少不是应该更快吗后来用性能计数器一看问题出在Cache Miss。表太大了每个桶平均访问间隔很远。一个key做线性探测第一次访问的缓存行里可能只有一个有效的slot剩下都是被浪费的预取。而把负载因子从0.25提高到0.7之后表变小了同一块缓存能覆盖的slot数量多了一倍多命中率立刻上来。这个坑的启示是哈希性能不只是冲突数量决定的还有内存层级。别盲目追求低负载因子要实际测一下不同容量下的耗时曲线。我当时在负载因子0.8附近拿到了最低延迟。4.3 坑三扩容时的并发读写竞争业务一开始没想好表大小运行到一半发现活跃ID超过预期需要扩容。如果只是单线程扩容很简单申请新表重新哈希所有key释放旧表。但在算子并行环境里其他核可能正在读旧表另一个线程在写新表指针切换的一瞬间读请求可能拿到一半的旧数据。我最终的规避方案是静态分配足够大的表不支持运行时扩容。这个限制对固定规模模型是能接受的。如果将来要做容量弹性我会用版本号双缓冲读方通过版本号判断当前该读哪张表写方把新表构建完成后再发布新版本。但这对算子的同步机制要求很高普通推理场景不建议轻易上。4.4 坑四多核竞争导致“看不见”的偶发错误哈希算子单测通过、batch推理偶发错一次、再跑又好了这种问题最折磨人。后来我总结出一个固定套路构造一组“特意撞桶”的测试数据保证多个key映射到同一个桶和同一条缓存行然后在单核模式下跑出一个黄金结果再切多核模式逐段对比。还有一个好用的手段是给每次CAS打印详细的key、slot、失败原因。虽然日志打印本身会拖慢很多倍但在定位阶段完全值得。几乎每个并发哈希算子的问题最终都能靠这种“缩小并发规模人为制造冲突”的方式复现出来。5. 接进模型后的收益与后续可扩展方向5.1 端到端收益把哈希算子放进原先模型流程后最大的变化不是把某个替代函数跑得快了一点而是整条链路少了几次跨设备搬运。原来CPU查完得到行号再拷贝给加速卡现在ID直接进卡卡内完成映射和Embedding读取数据形状也更像一个完整的端到端算子。实际项目里的CTR模型推理p95延迟从6.8ms降到了5.1ms其中哈希算子贡献了大约1.2ms的下降其余来自减少了拷贝和排队。Embedding查询的吞吐也从每秒不到300万ID提升到了每秒近2000万ID。这个提升对高并发场景很值。5.2 后续方向融合、稀疏更新与多级哈希接下来有几个方向我觉得很有意思。一个是把哈希查找和Embedding读取融合成一个算子。现在的设计还分两步哈希表输出行号下一个算子再根据行号取向量。如果能融合行号可以完全留在片上连输出数组都省掉省一次全局内存读写对于大批量场景非常可观。另一个方向是支持稀疏在线更新。线上模型如果要做增量学习就需要对Embedding表实时更新每来一个样本某些ID对应的向量要原地做SGD更新。这时候哈希算子不只是查表和插入还要支持“读取向量-更新-写回”的原子流程。这对value的并发一致性要求更高但也更有价值。还有一个思路是多级哈希。热点ID只占很小比例但访问量巨大可以给热点ID单独开一个小型高速哈希表冷门ID走大表。配合count-min sketch或者布隆过滤器能进一步把热点查找的缓存命中率拉高。这个方向对超大Embedding场景特别友好。5.3 给后来者的一点个人体会最后分享一个我自己反复用的小习惯在写这类哈希算子之前先把哈希函数在不同key分布下的输出打出来看一眼。很多看似复杂的并发问题根子其实在初始哈希分布不够均匀。调好哈希函数、定好负载因子、再谈原子操作和内存布局整个开发会顺很多。如果你也在做类似的稀疏模型加速我建议别一上来就追高级数据结构。先把最简单的线性探测哈希放到目标硬件上跑一遍用性能计数器看cache miss和原子冲突率再决定要不要引入更复杂的Robin Hood或者多级表。数据驱动的优化比拍脑袋改结构要靠谱得多。