LLM-DataDist APIs

LLM-DataDist: LLM distributed cluster and data acceleration component, which provides the cluster key-value (KV) data management to support separate deployment of full graphs and incremental graphs.

The following products are supported:
- Atlas A2 training product/Atlas A2 inference product: 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.
- Atlas A3 training product/Atlas A3 inference product

Currently, Python 3.9, Python 3.11, and Python 3.12 are supported. For details about how to install it, visit the Python official website: https://www.python.org/.
A maximum of 50 GB device memory can be registered. The larger the registered memory is, the more the OS memory is occupied.

The LLM-DataDist APIs are listed as follows.

LLM-DataDist

**Table 1** LLM-DataDist APIs
API	Description
LLMDataDist Constructor	Constructs LLMDataDist.
init	Initializes an LLM-DataDist instance.
finalize	Releases an LLM-DataDist instance.
link_clusters	Establishes a link.
unlink_clusters	Disconnects a link.
check_link_status	Quickly checks whether the link status is normal.
kv_cache_manager	Obtains the KvCacheManager instance.
switch_role	Switches the role of the current LLM-DataDist instance. You are advised to use this function only when PagedAttention is enabled.

LLMConfig

**Table 2** LLMConfig APIs
API	Description
LLMConfig Constructor	Constructs an LLMConfig instance.
generate_options	Generates a configuration item dictionary.
device_id	Sets the device ID of the current process, corresponding to the underlying ge.exec.deviceId configuration item.
sync_kv_timeout	Sets the timeout interval for APIs (such as the KV pull API), corresponding to the underlying llm.SyncKvCacheWaitTime configuration item.
enable_switch_role	Enables or disables smooth role switchover, corresponding to the underlying llm.EnableSwitchRole configuration item.
ge_options	Configures additional GE configuration items.
listen_ip_info	Sets the cluster listening information on the PROMPT side, corresponding to the underlying llm.listenIpInfo configuration item.
mem_utilization	Configures the utilization ratio of the memory specified by ge.flowGraphMemMaxSize. Default value: 0.95
buf_pool_cfg	Specifies the memory level configuration to improve memory allocation performance and utilization.

KvCacheManager

**Table 3** KvCacheManager APIs
API	Description
KvCacheManager Constructor	Constructs a KvCacheManager instance.
is_initialized	Checks whether the KvCacheManager instance has been initialized.
allocate_cache	Allocates a cache. After a cache is successfully allocated, it is referenced by cache_id and cache_keys at the same time. The resources occupied by the cache are released only after these references are removed.
deallocate_cache	Deallocates a cache.
remove_cache_key	Removes CacheKey. This API can be called only when LLMRole is set to PROMPT.
pull_cache	Pulls the KV cache from the corresponding Prompt node to the local KV cache based on CacheKey. This API can be called only when LLMRole is set to DECODER.
copy_cache	Copies the KV.
get_cache_tensors	Obtains cache tensors.
allocate_blocks_cache	Allocates the cache for multiple blocks in the PagedAttention scenario.
pull_blocks	Pulls the KV cache from the corresponding Prompt node to the local KV cache via a block list based on BlocksCacheKey in the PagedAttention scenario. This API can be called only when LLMRole is set to DECODER.
copy_blocks	Copies the KV in the PagedAttention scenario.
swap_blocks	Swaps in and out the CPU cache (cpu_cache) and NPU cache (npu_cache).
transfer_cache_async	Asynchronously transfers the KV cache in a hierarchical manner.

KvCache

**Table 4** KVCache APIs
API	Description
KvCache Constructor	Constructs a KV cache.
cache_id	Obtains the KV cache ID.
cache_desc	Obtains the KV cache description.
per_device_tensor_addrs	Obtains the KV cache address.
create_cpu_cache	Creates a CPU cache.

LLMClusterInfo

**Table 5** LLMClusterInfo APIs
API	Description
LLMClusterInfo Constructor	Constructs an LLMClusterInfo instance.
remote_cluster_id	Sets the peer cluster ID.
append_local_ip_info	Adds the IP address of the local cluster.
append_remote_ip_info	Adds the IP address of the remote cluster.

CacheTask

**Table 6** CacheTask APIs
API	Description
CacheTask Constructor	Constructs a CacheTask instance.
synchronize	Waits until the transmission at all layers is complete and obtains the overall execution result.
get_results	Waits until the transmission at all layers is complete and obtains the execution result of each TransferConfig instance.

Others

**Table 7** Other APIs
API	Description
LLMRole	Enumerates values of LLMRole.
Placement	Represents the device type where the cache is located. It is a field in CacheDesc.
CacheDesc	Constructs a CacheDesc instance.
CacheKey	Constructs a CacheKey instance.
CacheKeyByIdAndIndex	Constructs a CacheKeyByIdAndIndex instance, which is usually used as the parameter type in pull_cache.
BlocksCacheKey	Constructs a BlocksCacheKey instance in the PagedAttention scenario.
LayerSynchronizer	Waits for the specified layer of the model to complete execution. Users need to inherit from LayerSynchronizer and implement this API. This API is called during the execution of KvCacheManager.transfer_cache_async. When the API returns a success status, the cache transmission for the current layer begins.
TransferConfig	Constructs a TransferConfig instance.
TransferWithCacheKeyConfig	Constructs a TransferWithCacheKeyConfig instance.
LLMException	Obtains the error code of an exception. For details about the error code list, see LLMStatusCode.
LLMStatusCode	Enumerates values of LLMStatusCode.
DataType	Enumerates values of DataType.

Parent topic: API Reference