pull_blocks

函数功能

PagedAttention场景下，拉取kv。

函数原型

pull_blocks(llm_req: LLMReq, prompt_blocks: List[int] = [], decoder_blocks: List[int] = [])

参数说明

参数名称	数据类型	取值说明
llm_req	LLMReq	请求信息。
prompt_blocks	List[int]	prompt的block index列表。
decoder_blocks	List[int]	decoder的block index列表。

调用示例

from llm_datadist import LLMDataDist, LLMRole, ModelConfig, ModelRunner, LLMReq
llm_datadist = LLMDataDist(LLMRole.DECODER, 0)
model_config = ModelConfig()
kv_nums = 80
model_config.kv_shapes = ["1,1024" for _ in range(kv_nums)]
model_config.kv_dtypes = [DataType.DT_FLOAT for _ in range(kv_nums)]
model_options = model_config.generate_options()
class TestModelRunner(ModelRunner):    
   def run_model(self, kv_cache: KvCache, input_tensors, **kwargs):        
       return []
llm_model = llm_datadist.add_model(model_options, TestModelRunner())
llm_req = LLMReq()
llm_model.pull_blocks(llm_req, [], [2, 3])

返回值

正常情况下无返回值。

异常情况会抛出LLMException。

约束说明

无

父主题： LLMModel