简介
为了方便后续的描述,我们对于不同的推理框架抽象为几个模块:
- 资源初始化模块
- KV Cache管理模块。包括KV Cache的内存的创建,分配(PA场景)以及销毁。
- 模型推理模块。
- 资源释放模块。
本章节主要是介绍开发者如何在推理框架中使能LLM-DataDist的能力,LLM-DataDist按照业务支持场景的差异提供了两套接口。
仅支持D2D传输的接口开发流程
- 找到推理框架中的资源初始化模块,在该阶段中调用LLM-DataDist的初始化接口和建链接口。
- 找到推理框架中的KV Cache管理模块,在该阶段中调用LLM-DataDist的KV Cache申请接口,将申请好的KV Cache转化为不同框架的KV Cache类型,在原位置直接替换。
- 推理框架要能够拆分出Prefill阶段和Decoder阶段,对推理脚本进行分离部署,部署到不同的集群节点上。
- 在Decoder阶段执行前需要接受来自Prefill阶段的输出作为输入,同时调用LLM-DataDist提供的KV Cache传输的接口拉取Prefill侧缓存的KV Cache。
- 分别执行Prefill推理脚本和Decoder推理脚本。
- 在框架资源释放模块释放llm_datadist相关资源。
支持D2D/D2H/H2D的接口开发流程
这套接口新增支持了D2H和H2D场景,还额外提供了将用户内存注册给llm_datadist的使用方式,同时会淡化角色这个概念,接口调用不受角色限制。
通过llm_datadist申请管理cache内存的开发流程和上述仅支持D2D传输的接口开发流程一致
将用户注册给llm_datadist的接口开发流程:
- 找到推理框架中的资源初始化模块,在该阶段中调用LLM-DataDist的初始化接口。
- 找到推理框架中的KV Cache管理模块,将框架中申请好的cache的内存地址注册给llm_datadist。
- 在框架模型推理前调用查询注册内存状态的接口,等待用户内存注册完成。
- 推理框架要能够拆分出Prefill阶段和Decoder阶段,对推理脚本进行分离部署,部署到不同的集群节点上。
- 根据业务上需要进行数据传输的节点,调用llm_datadist提供的数据传输接口将数据拉到对应的节点上。
- 分别执行Prefill推理脚本和Decoder推理脚本。
- 在框架资源释放模块释放llm_datadist相关资源。