check_capacity

产品支持情况

产品

是否支持

Atlas 800I A2 推理产品/A200I A2 Box 异构组件

√

Atlas A2 训练系列产品

x

Atlas 200I/500 A2 推理产品

x

Atlas 推理系列产品

x

Atlas 训练系列产品

x

函数功能

在KvCacheManager模式下，在全量predict接口执行前调用，用于在推理前检查本轮想要执行请求的KV Cache是否可以缓存下。调用complete_request释放KV Cache占用。

函数原型

1
check_capacity(seq_len: int) -> CapacityState

参数说明

参数名称

数据类型

取值说明

seq_len

int

请求token数量。

  • 单batch场景是从LLMReq的prompt_length获取值。
  • 在多batch场景下，seq_len=max_prompt_len_per_batch * batch_size。

调用示例

1
2
3
4
5
6
from llm_datadist import LLMDataDist, LLMRole
llm_datadist = LLMDataDist(LLMRole.PROMPT, 0)
...
llm_req = LLMReq()
llm_req.prompt_length = 1
capacity_state = llm_datadist.check_capacity(llm_req.prompt_length)

返回值

正常情况下返回CapacityState。

参数错误可能抛出TypeError或ValueError。

约束说明

父主题： LLMDataDist