Add a inplace concat custom op based on CUDA VMM API (resubmitted)

[lszxb]

PR types

Performance optimization

PR changes

Others

Description

这一PR尝试为当前的大模型推理过程增加基于CUDA VMM API的inplace concat支持（原理类似于vAttention），从而避免在每一个解码步都复制一次整个KV Cache。
该功能暂时只实现了自定义算子，未来还需要增加相关的pass以自动适配其他模型。
目前这一PR在llama模型上应用了这一方案，在3072 input+1024 output的情况下大约有10%的提升。

目前主要的思路是：

• 使用一种特殊的Tensor，其显存由VMM API分配，这种Tensor使用特殊的phi::Allocation，在创建时预留大量的虚拟地址空间，可以在必要时分配物理页映射到虚拟地址空间。
• 为了兼容剩余的调用，cache的shape为batch x seq_len x num_head x head_dim，但由于状态在cache的尾部追加，cache的内存布局应该是seq_len x batch x num_head x head_dim。
  vtensor_reserve_one_token自定义算子的语义大致如下：
• 如果key_cache不是VTensor，则新分配一个VTensor，并将原先key_cache中的数据复制到这个新的VTensor中。然后使用VTensor的扩展机制，在尾部预留新的一个token的空间，并将key_states复制到这个新的空间中。
• 如果key_cache是VTensor，直接使用VTensor的扩展机制，在尾部预留新的一个token的空间，并将key_states复制到这个新的空间中。

目前可能存在的问题：

• 仅支持每次追加1个token的空间。
• 目前分配的虚拟地址空间大小和block大小为定值（1GiB与32MiB），可能暴露相关的API给用户进行调整会更好？
• 输入和输出的key_cache共享同一块空间，在某些情况下可能会产生冲突。
• 该方法依赖于每个step使用的kv cache是同一个Tensor，若有某些其他操作改变了kv cache的Tensor（比如说clone到另一个Tensor），则会导致失效，因此也需要配合这个PR的优化才可使用（assign_out_操作会导致复制）。
• 通过该算子分配的显存无法使用现有的Allocator进行统一管理。

本PR还包括了以下两个PR的内容：

• [Performance] Remove the redundant pd_op.assign_out_ op at the end of while loop #9002
• [Performance] Avoid cuda sync in postprocess of LLM decoding #9011

[paddle-bot]

Thanks for your contribution!

[yuanlehome]

新增的几个文件开头放一下版权声明

[lszxb]

加上了

[yuanlehome]

类似注释都可以给删掉

[yuanlehome]

这个文件不需要更改，先暂时只在gpu下使用就行

[lszxb]

好的，这几行删掉了

[yuanlehome]

Suggested change

	if is_paddlenlp_ops_available():
	import paddlenlp_ops
	inference_config.enable_custom_passes([
	"remove_assign_out_pass", # remove the assign_out_ op at the end of while loop
	"apply_vtensor_concat_pass", # replace concat op with vtensor implementation
	])
	try:
	import remove_assign_out_pass, apply_vtensor_concat_pass from paddlenlp_ops
	inference_config.enable_custom_passes([
	"remove_assign_out_pass", # remove the assign_out_ op at the end of while loop
	"apply_vtensor_concat_pass", # replace concat op with vtensor implementation
	])
	except:
	pass

这样修改吧

[lszxb]

paddlenlp_ops里没有pass的对象，我换成了新加的算子vtensor_reserve_one_token

[codecov]

Codecov Report

Attention: Patch coverage is 28.57143% with 5 lines in your changes missing coverage. Please review.

Project coverage is 52.24%. Comparing base (81f5ab5) to head (d5a9393).
Report is 2 commits behind head on develop.

Files with missing lines	Patch %	Lines
paddlenlp/generation/utils.py	0.00%	4 Missing ⚠️
paddlenlp/generation/logits_process.py	66.66%	1 Missing ⚠️

Additional details and impacted files

@@             Coverage Diff             @@
##           develop    #9320      +/-   ##
===========================================
- Coverage    52.92%   52.24%   -0.69%     
===========================================
  Files          661      671      +10     
  Lines       107069   109655    +2586     
===========================================
+ Hits         56670    57288     +618     
- Misses       50399    52367    +1968     

☔ View full report in Codecov by Sentry.
📢 Have feedback on the report? Share it here.