FlashInfer项目中Hopper架构下非连续分页KV存储的预填充注意力机制实现

背景介绍

FlashInfer是一个高性能的注意力机制实现库，特别针对不同NVIDIA GPU架构进行了优化。在最新版本中，项目引入了对非连续分页KV存储(paged_kv)的支持，允许用户使用自定义步长(stride)的内存布局。这种灵活性对于某些特定场景下的内存优化尤为重要。

技术挑战

在Hopper架构(SM90)之前的GPU上，预填充注意力(prefill attention)内核已经能够很好地支持自定义步长的分页KV存储布局，例如[max_num_pages, num_layer, 2, page_size, num_heads, head_dim]这样的内存结构。然而，当尝试在Hopper架构上使用相同配置时，开发者发现相关参数定义中缺少对自定义步长的直接支持。

解决方案

FlashInfer团队针对Hopper架构采用了不同的实现策略：

1. 预处理偏移量：不同于直接在内核中计算指针偏移，团队选择在调用内核前预先计算好所有偏移量。这种设计决策源于性能考量——在内核中进行指针算术运算会导致生产者性能下降。

2. 偏移量转换：将分页存储的步长(stride_page)转换为具体的偏移量，前提是满足stride_block % stride_n == 0的条件。这种转换确保了内存访问的高效性。

3. 性能优化考量：通过尽量减少内核中的指针运算，预计算偏移量的方法显著提升了整体性能，特别是在处理非连续内存布局时。

实现细节

在代码实现层面，FlashInfer通过以下方式支持Hopper架构的非连续分页KV存储：

• 在Python接口层进行步长到偏移量的转换
• 将转换后的偏移量传递给底层CUDA内核
• 在内核中使用预处理好的偏移量进行高效内存访问

未来优化方向

虽然当前实现已经解决了基本功能需求，但仍有优化空间：

1. 统一接口设计：考虑将偏移量计算逻辑整合到PrefillSM90Plan中，使SM90前和SM90后的接口使用方式保持一致。

2. 条件性优化：如果不同层使用相同的索引模式，可以将偏移量计算逻辑上移到计划(plan)阶段，进一步减少重复计算。

总结

FlashInfer项目通过创新的预处理偏移量方法，成功在Hopper架构上实现了对非连续分页KV存储的支持。这种设计不仅解决了功能需求，还通过减少内核中的指针运算显著提升了性能。随着项目的持续发展，接口统一化和进一步性能优化将是重点方向。