TransformerEngine中Flash Attention的qkvpacked格式使用指南

理解qkvpacked格式

在TransformerEngine项目中，Flash Attention提供了一种高效的注意力机制实现方式。其中flash_attn_varlen_qkvpacked_func函数接受一个特殊的qkv张量格式，该张量将查询(Q)、键(K)和值(V)三个矩阵打包在一起，形状为(total, 3, nheads, headdim)，其中total表示批次中所有token的总数。

迁移到DotProductAttention类

当从底层函数迁移到更高级的DotProductAttention类时，需要注意几个关键点：

1. qkv_format的选择

正确的格式应该是thd，其中：

• t：批次中所有token的总数（注意文档中之前的描述有误）
• h：注意力头的数量
• d：每个头的维度

2. 注意力掩码设置

对于非因果注意力(non-causal attention)，应使用attn_mask_type=padding。即使不需要实际的注意力掩码，也需要通过cu_seqlens_q和cu_seqlens_kv参数提供序列长度信息，这有助于后端了解每个序列的实际长度。

3. 序列长度参数

cu_seqlens_q和cu_seqlens_kv应该是累积序列长度。例如，对于一个包含3个序列的批次，长度分别为3、2和4，则累积序列长度应为[0, 3, 5, 9]。

如果序列之间没有填充（即token是连续排列的，没有PAD token），则不需要指定cu_seqlens_q_padded或cu_seqlens_kv_padded参数。

实际应用建议

在实际应用中，建议遵循以下最佳实践：

1. 确保输入张量的形状严格符合thd格式要求
2. 即使不需要掩码，也要提供序列长度信息
3. 对于非填充序列，可以省略_padded参数
4. 考虑使用项目提供的单元测试作为参考实现

通过正确配置这些参数，可以充分利用TransformerEngine中Flash Attention的高效实现，同时保持代码的清晰性和可维护性。