背景
大模型计算中，往往会引入Padding操作处理不等长序列，但是这些Padding的token会引发冗余计算的问题，因此需要把Padding的内容丢掉，直接使用有效的token进行attention的计算。在Paddle中，是使用get_padding_offset操作完成这一步骤的。get_padding_offset函数的定义在这里。
def get_padding_offset(bsz, max_seq_len, seq_lens_this_time):
# 计算每个序列需要填充的长度，并累加得到累计偏移量
cum_offsets_now = paddle.cumsum(max_seq_len - seq_lens_this_time)
# 创建累计偏移量张量，长度为bsz+1,第一个元素为0
cum_offsets = paddle.zeros(shape=(bsz + 1), dtype="int32")
cum_offsets[1:] = cum_offsets_now # 将累加结果放入cum_offsets[1:]
# 计算所有序列的实际token总数（不包含填充部分）
token_num = paddle.sum(seq_lens_this_time)
# 初始化填充偏移量张量，长度为所有实际token数
padding_offsets = paddle.zeros(shape=(token_num), dtype="int32")
# 初始化查询Q和键K的累计序列长度张量
cu_seqlens_q = paddle.zeros(shape=(bsz + 1), dtype="int32")
cu_seqlens_k = paddle.zeros(shape=(bsz + 1), dtype="int32")
# 遍历批次中的每个序列
for i in range(bsz):
seq_len_now = seq_lens_this_time[i] # 当前序列的实际长度
cum_offset = cum_offsets[i] # 当前序列的累计偏移量
# 遍历当前序列中的每个token
for j in range(seq_len_now):
# 计算该token在预填充后序列中的位置，并设置其预填充偏移量
padding_offsets[i * max_seq_len + j - cum_offset] = cum_offset
# 计算累计序列长度（用于高效实现注意力机制）
cum_seq_len = (i + 1) * max_seq_len - cum_offsets[i + 1]
cu_seqlens_q[i + 1] = cum_seq_len
cu_seqlens_k[i + 1] = cum_seq_len
return padding_offsets, cum_offsets[:-1], cu_seqlens_q, cu_seqlens_k
get_padding_offset说明
作用：get_padding_offset主要用于批处理中不同序列长度的填充问题，生成注意力机制中所需要的偏移量和累计序列长度信息。在Transformer模型的优化实现中，特别是在处理不等长序列时经常会用到类似的方法。
参数说明：
- bsz: 批次大小
- max_seq_len: 序列的最大长度
- seq_lens_this_time: 当前批次中每个序列的实际长度（Tensor)
核心功能：生成四个关键输出
- padding_offsets: 每个token在填充后的序列中的偏移量
- cum_offsets: 每个序列的累计偏移量
- cu_seqlens_q: 查询矩阵Q的累计序列长度
- cu_seqlens_k: 键矩阵K的累计序列长度
数学原理
这个函数实现了一种高效处理不等长序列的方法，主要用于优化Transformer模型中的注意力计算。在处理不等长序列时，通常有两种策略
- 填充：将所有序列填充到长度相同，会导致计算冗余
- 压缩：将所有有效token排列到一起，避免计算冗余 这个函数实现了压缩策略，通过生成padding_offsets和cu_seqlens张量，可以
- 高效的将填充后的张量转换为压缩形式
- 在注意力机制时，快速定位每个token在原始序列中的位置
- 避免在预填充位置上进行无用的计算，提高内存和计算效率
具体解释
假设bsz=3, max_seq_len=5, seq_lens_this_time=[3,4,2]
1、首先计算便宜量
cum_offsets_now = paddle.cumsum(max_seq_len - seq_lens_this_time)
# 填充长度 = [5-3, 5-4, 5-2] = [2, 1, 3]
# 累积填充长度 = [2, 2+1, 2+1+3] = [2, 3, 6]
2、初始化累计偏移量数组
cum_offsets = [0, 2, 3, 6] # 长度为bsz+1=4
3、计算token总数
token_num = sum([3, 4, 2]) = 9
4、填充偏移量和累计序列长度计算
序列1：（长度=3）
- 累计偏移量cum_offset=cum_offsets[0] = 0
- 循环填充
padding_offsets和cu_seqlens
| j | i*max_seq_len - cum_offset + j | padding_offsets 值 | cum_seq_len | cu_seqlens_q/k[1] |
|---|---|---|---|---|
| 0 | 0*5 - 0 + 0 = 0 | 0 | 5-2=3 | 3 |
| 1 | 0*5 - 0 + 1 = 1 | 0 | 5-2=3 | 3 |
| 2 | 0*5 - 0 + 2 = 2 | 0 | 5-2=3 | 3 |
序列2：（长度=4）
- 累计偏移量cum_offset = cum_offsets[1] = 2
- 循环填充
padding_offsets和cu_seqlens
| j | i*max_seq_len - cum_offset + j | padding_offsets 值 | cum_seq_len | cu_seqlens_q/k[2] |
|---|---|---|---|---|
| 0 | 1*5 - 2 + 0 = 3 | 2 | 10-3=7 | 7 |
| 1 | 1*5 - 2 + 1 = 4 | 2 | 10-3=7 | 7 |
| 2 | 1*5 - 2 + 2 = 5 | 2 | 10-3=7 | 7 |
| 3 | 1*5 - 2 + 3 = 6 | 2 | 10-3=7 | 7 |
序列 3 (长度 = 2)
- 累积偏移量
cum_offset = cum_offsets[2] = 3 - 循环填充
padding_offsets和cu_seqlens
| j | i*max_seq_len - cum_offset + j | padding_offsets 值 | cum_seq_len | cu_seqlens_q/k[3] |
|---|---|---|---|---|
| 0 | 2*5 - 3 + 0 = 7 | 3 | 15-6=9 | 9 |
| 1 | 2*5 - 3 + 1 = 8 | 3 | 15-6=9 | 9 |
5、最终结果
padding_offsets = [0, 0, 0, 2, 2, 2, 2, 3, 3]
cum_offsets = [0, 2, 3] # 去掉最后一个元素
cu_seqlens_q = [0, 3, 7, 9]
cu_seqlens_k = [0, 3, 7, 9]
实际使用中的意义
假设我们有三个填充后的序列
seq1: [a1, a2, a3, 0, 0]
seq2: [b1, b2, b3, b4, 0]
seq3: [c1, c2, 0, 0, 0]
填充后的矩阵形状是[3,5], 但实际有效token只有9个，使用我们计算padding_offsets和cu_seqlens, 我们可以
1、将填充后的矩阵压缩为[9]的一维数组
compressed = [a1, a2, a3, b1, b2, b3, b4, c1, c2]
2、在计算注意力时，快速定位每个token的位置
- padding_offsets告诉我们每个token在原始填充序列的偏移量
- cu_seqlens告诉我们每个序列在压缩数组中的起始位置
这种方法避免了在填充位置上进行无用注意力计算，显著提高了计算效率，特别是在处理大量不等长序列时。
