allocate_blocks_cache

Applicability

Product	Supported (Yes/No)
Atlas 350 Accelerator Card	No
Atlas A3 training product/Atlas A3 inference product	Yes
Atlas A2 training product/Atlas A2 inference product	Yes
Atlas 200I/500 A2 inference product	No
Atlas inference product	No
Atlas training product	No

Note: For the Atlas A2 training product/Atlas A2 inference product, only the Atlas 800I A2 inference server and A200I A2 Box heterogeneous subrack are supported.

Function Description

In the PagedAttention scenario, allocates the cache with multiple blocks. After the cache is successfully allocated, the memory can be released through deallocate_cache.

Prototype

allocate_blocks_cache(cache_desc: CacheDesc, blocks_cache_key: Optional[BlocksCacheKey] = None)

Parameters

Parameter	Data Type	Description
cache_desc	CacheDesc	Cache description.
blocks_cache_key	Optional[BlocksCacheKey]	This parameter can be set only when LLMRole is set to PROMPT, and it is used for the Decode side to pull KV cache.

Example

from llm_datadist import BlocksCacheKey
num_blocks = 1000
cache_desc = CacheDesc(80, [num_blocks , 128 * 1024], DataType.DT_FLOAT16)
blocks_cache_key = BlocksCacheKey(0, 0)
kv_cache = kv_cache_manager.allocate_blocks_cache(cache_desc, blocks_cache_key)

Returns

In normal cases, a KvCache instance is returned.

If a parameter is incorrect, a TypeError or ValueError may be thrown.

If the execution time exceeds the value of sync_kv_timeout, an LLMException is thrown.

Constraints

This API does not support concurrent calls.

Parent topic: KvCacheManager