(bsz, max_seq_len, seq_lens_this_time)
| 334 | |
| 335 | |
| 336 | def get_padding_offset(bsz, max_seq_len, seq_lens_this_time): |
| 337 | cum_offsets_now = paddle.cumsum(max_seq_len - seq_lens_this_time, dtype="int32") |
| 338 | cum_offsets = paddle.zeros(shape=(bsz + 1), dtype="int32") |
| 339 | cum_offsets[1:] = cum_offsets_now |
| 340 | token_num = paddle.sum(seq_lens_this_time) |
| 341 | padding_offsets = paddle.zeros(shape=(token_num), dtype="int32") |
| 342 | batch_id_per_token = paddle.zeros(shape=(token_num), dtype="int32") |
| 343 | cu_seqlens_q = paddle.zeros(shape=(bsz + 1), dtype="int32") |
| 344 | cu_seqlens_k = paddle.zeros(shape=(bsz + 1), dtype="int32") |
| 345 | for i in range(bsz): |
| 346 | seq_len_now = seq_lens_this_time[i] |
| 347 | cum_offset = cum_offsets[i] |
| 348 | for j in range(seq_len_now): |
| 349 | padding_offsets[i * max_seq_len - cum_offset + j] = cum_offset |
| 350 | batch_id_per_token[i * max_seq_len - cum_offset + j] = i |
| 351 | cum_seq_len = (i + 1) * max_seq_len - cum_offsets[i + 1] |
| 352 | cu_seqlens_q[i + 1] = cum_seq_len |
| 353 | cu_seqlens_k[i + 1] = cum_seq_len |
| 354 | if fastdeploy.platforms.current_platform.is_cuda(): |
| 355 | return batch_id_per_token, cum_offsets[:-1], cu_seqlens_q, cu_seqlens_k |
| 356 | else: |
| 357 | return padding_offsets, cum_offsets[:-1], cu_seqlens_q, cu_seqlens_k |
| 358 | |
| 359 | |
| 360 | def remove_padding(seq_lens, cu_seq_lens, inputs, token_num): |
no test coverage detected