英伟达megatron的数据layout为啥要用sbh（先seq）而非bsh（先batch）?
2023-09-10 阅读 6
因为 Transformer 的关键计算（注意力的 softmax、Q·K^T、KV 缓存等）是以序列维度为主的，序列维度放在最前面能保证内存连续性和更好的缓存局部性，减少转置和不必要的 stride，便于用高效的 fused kernel/ tensor core 实现，也方便做序列切分的并行（sequence parallel）和推理时对 KV cache 的索引管理，整体吞吐和延迟更优，所以常用 sbh 而不是 bsh。
更新于 2025年12月26日
