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 products /Atlas A2 inference products : For theAtlas A2 training products /Atlas A2 inference products , only the Atlas 800I A2 inference server and A200I A2 Box heterogeneous subrack are supported.Atlas A3 training products /Atlas A3 inference products
- Currently, Python 3.9 to 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
API |
Description |
|---|---|
Constructs LLMDataDist. |
|
Initializes an LLM-DataDist instance. |
|
Releases the LLMDataDist. |
|
Establishes a link. |
|
Disconnects a link. |
|
Quickly checks whether the link status is normal. |
|
Obtains the KvCacheManager instance. |
|
Switches the role of the current LLM-DataDist instance. You are advised to use this function only when PagedAttention is enabled. |
LLMConfig
API |
Description |
|---|---|
Constructs LLMConfig. |
|
Generates a configuration item dictionary. |
|
Sets the device ID of the current process, corresponding to the underlying ge.exec.deviceId configuration item. |
|
Sets the timeout interval for APIs (such as the KV pull API), corresponding to the underlying llm.SyncKvCacheWaitTime configuration item. |
|
Enables or disables the smooth role switchover, corresponding to the llm.EnableSwitchRole configuration item at the bottom layer. |
|
Configures additional GE configuration items. |
|
Sets the cluster listening information on the PROMPT side, which corresponds to the llm.listenIpInfo configuration item at the bottom layer. |
|
Configures the memory usage of ge.flowGraphMemMaxSize. Default value: 0.95 |
|
Specifies the memory level configuration to improve memory allocation performance and utilization. |
KvCacheManager
API |
Description |
|---|---|
Constructs a KvCacheManager instance. |
|
Checks whether the KvCacheManager instance has been initialized. |
|
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. |
|
Deallocates a cache. |
|
Removes CacheKey. This API can be called only when LLMRole is set to PROMPT. |
|
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. |
|
Copies the KV. |
|
Obtains a cache tensor. |
|
Allocates the cache for multiple blocks in the PagedAttention scenario. |
|
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. |
|
Copies the KV in the PagedAttention scenario. |
|
Swaps in and out the CPU cache (cpu_cache) and NPU cache (npu_cache). |
|
Asynchronously transfers the KV cache in a hierarchical manner. |
KvCache
API |
Description |
|---|---|
Constructs a KV cache. |
|
Obtains the KV cache ID. |
|
Obtains the KV cache description. |
|
Obtains the KV cache address. |
|
Creates a CPU cache. |
LLMClusterInfo
API |
Description |
|---|---|
Constructs LLMClusterInfo. |
|
Sets the peer cluster ID. |
|
Adds the IP address of the local cluster. |
|
Adds the IP address of the remote cluster. |
CacheTask
API |
Description |
|---|---|
Constructs a CacheTask. |
|
Waits until the transmission at all layers is complete and obtains the overall execution result. |
|
Waits until the transmission at all layers is complete and obtains the execution result of each TransferConfig instance. |
Others
API |
Description |
|---|---|
Enumerates values of LLMRole. |
|
Represents the device type where the cache is located. It is a field in CacheDesc. |
|
Constructs CacheDesc. |
|
Constructs CacheKey. |
|
Constructs a CacheKeyByIdAndIndex instance, which is usually used as the parameter type in pull_cache. |
|
Constructs a BlocksCacheKey instance in the PagedAttention scenario. |
|
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 when KvCacheManager.transfer_cache_async is executed. When a success message is returned, the cache transfer at the current layer starts. |
|
Constructs a TransferConfig instance. |
|
Constructs a TransferWithCacheKeyConfig instance. |
|
Obtains the error code of an exception. For details about the error code list, see LLMStatusCode. |
|
Enumerates values of LLMStatusCode. |
|
Enumerates values of DataType. |