LLM-DataDist APIs

LLM-DataDist is an LLM distributed cluster and data acceleration component. It provides clustered KV data management capabilities and supports disaggregated deployment of full graphs and incremental graphs.

The following products are supported:
- Atlas A2 training product/Atlas A2 inference product: For Atlas A2 training product/Atlas A2 inference product, only the Atlas 800I A2 inference server and A200I A2 Box heterogeneous subrack are supported. When the HCCS transmission protocol is used in the server, only D2D transmission is supported.
- Atlas A3 training product/Atlas A3 inference product: When the HCCS transmission protocol is used, the host memory cannot be used as the remote cache.

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 and 20 GB host memory can be registered. Larger registered memory results in higher OS memory usage. This restriction applies to the following products:
- Atlas A3 training product/Atlas A3 inference product
- Atlas A2 training product/Atlas A2 inference product: For Atlas A2 training product/Atlas A2 inference product, only the Atlas 800I A2 inference server and A200I A2 Box heterogeneous subrack are supported.

The LLM-DataDist APIs are as follows.

LLM-DataDist

**Table 1** APIs of the LLM-DataDist
API	Description
LLMDataDist Constructor	Constructs LLMDataDist.
init	Initializes the LLMDataDist.
finalize	Releases the LLMDataDist.
link_clusters	Establishes a link unilaterally. The request is initiated by the client. The server/client roles are independent of the prompt/decoder roles. Setting listen_ip_info on a node designates the node as the server (port listens for links).
unlink_clusters	Disconnects a link unilaterally. The request can be initiated by the client to notify the server of disconnection. Alternatively, both the client and server can initiate forcible disconnection to clean up only local links.
switch_role	Switches the role of the current LLM-DataDist instance. You can also configure switch_options to switch the role to client or server.
link	Establishes a link by establishing a communicator in bilateral link establishment scenarios.
unlink	Disconnects a link in bilateral link establishment scenarios.
query_register_mem_status	Queries the status of the registered memory in bilateral link establishment scenarios.
cache_manager	Obtains a CacheManager instance.

LLMConfig

**Table 2** APIs of the LLMConfig
API	Description
LLMConfig Constructor	Constructs LLMConfig.
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	Configures the timeout interval for APIs that pull the KV, corresponding to the llm.SyncKvCacheWaitTime configuration item at the bottom layer.
ge_options	Configures additional GE configuration items.
listen_ip_info	Sets the cluster listening information on the Prompt side, corresponding to the llm.listenIpInfo configuration item at the bottom layer.
enable_cache_manager	Enables or disables CacheManager, corresponding to the llm.EnableCacheManager configuration item at the bottom layer. Set this item to True.
mem_pool_cfg	When CacheManager is enabled, configures the memory pool-related settings.
host_mem_pool_cfg	Configures the host memory pool when CacheManager is enabled.
enable_remote_cache_accessible	When CacheManager is enabled, enables or disables direct access to the remote cache.
rdma_traffic_class	Sets the traffic class of the RDMA NIC when CacheManager is enabled.
rdma_service_level	Sets the service level of the RDMA NIC when CacheManager is enabled.
local_comm_res	Configures local communication resources when CacheManager is enabled.
link_total_time	Sets the total timeout interval for HCCL link establishment failures.
link_retry_count	Sets the number of retries upon HCCL link establishment failures.
transfer_backend	Sets the transmission backend used by LLM-DataDist.

CacheManager

**Table 3** APIs of the CacheManager
API	Description
CacheManager Constructor	Constructs a CacheManager instance.
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	Releases the cache. If the cache is associated with a cache key during allocation, the actual release is delayed until all cache keys are pulled or remove_cache_key is executed.
remove_cache_key	Removes CacheKey. After CacheKey is removed, the cache cannot be pulled by pull_cache.
pull_cache	Obtains the cache from the peer node to the local cache based on CacheKey.
copy_cache	Copies the cache.
allocate_blocks_cache	In the PagedAttention scenario, allocates the cache with multiple blocks. After the cache is successfully allocated, the memory can be released through deallocate_blocks_cache.
deallocate_blocks_cache	In the PagedAttention scenario, releases the cache allocated by allocate_blocks_cache.
pull_blocks	In the PagedAttention scenario, the cache is pulled from the peer node in block list based on BlocksCacheKey to the local cache.
copy_blocks	Copies blocks in the PagedAttention scenario.
swap_blocks	Swaps in and out the cpu_cache and npu_cache.
register_cache	Registers a self-allocated memory.
register_blocks_cache	In the PagedAttention scenario, this API is called to register a self-allocated memory.
transfer_cache_async	Transfers the data of cache in asynchronously hierarchical mode.
push_blocks	Pushes the cache from the local node to the remote cache in block list based on BlocksCacheKey in the PagedAttention scenario.
push_cache	Pushes the cache from the local node to the remote cache based on the cache key (CacheKey).
remap_registered_memory	During LLM inference, if a memory UCE occurs (that is, ACL_ERROR_RT_DEVICE_MEM_ERROR is returned), the upper-layer framework needs to first determine whether the faulty memory is a KV cache memory. If it is not a KV cache memory, refer to the torch_npu.npu.restart_device API description in Ascend Extension for PyTorch Custom API Reference to obtain and fix the incorrect virtual address of the memory UCE. If it is a KV cache memory, additionally call this API to restore the KV cache memory registered with the NIC. This API is reserved, and the function is not supported currently.
unregister_cache	Deregisters a self-allocated memory.

Cache

**Table 4** APIs of the Cache
API	Description
Cache Constructor	Constructs cache.
cache_id	Obtains the cache ID.
cache_desc	Obtains the cache description.
tensor_addrs	Obtains the cache address.
create_cpu_cache	Creates a CPU cache.

LLMClusterInfo

**Table 5** APIs of the LLMClusterInfo
API	Description
LLMClusterInfo Constructor	Constructs LLMClusterInfo.
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
API	Description
CacheTask Constructor	Constructs a CacheTask.
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** Others
API	Description
LLMRole	The enumerated values of LLMRole.
RegisterMemStatus	The enumerated values of RegisterMemStatus.
Placement	A field in CacheDesc, indicating the type of device where the cache is located.
CacheDesc	Constructs CacheDesc.
MemType	Indicates the location of the address. It is usually used as the type of the mem_type parameter of MemInfo.
MemInfo	Constructs a MemInfo instance, which is typically used as a parameter type in the remap_registered_memory API of CacheManager.
CacheKey	Constructs CacheKey.
CacheKeyByIdAndIndex	Constructs a CacheKeyByIdAndIndex instance, which is usually used as the parameter type in pull_cache and push_cache of CacheManager.
BlocksCacheKey	In the PagedAttention scenario, constructs BlocksCacheKey.
LayerSynchronizer	Waits until the execution of the specified layer of the model is complete. You need to inherit LayerSynchronizer and implement this API. This API is called during the execution of transfer_cache_async. When the API returns a success status, the cache transmission for the current layer begins.
TransferConfig	Constructs TransferConfig.
TransferWithCacheKeyConfig	Constructs a TransferWithCacheKeyConfig instance.
LLMException	Obtains the error code of an exception. For details about the error code list, see LLMStatusCode.
LLMStatusCode	The enumerated values of LLMStatusCode.
DataType	Enumerates values of DataType.

Parent topic: API Reference