LLM-DataDist Data Structures

LlmRole

Specifies the role of an LLM-DataDist node.

enum class LlmRole : int32_t {
  kPrompt = 1,      // Role: Prompt
  kDecoder = 2,     // Role: Decoder
  kMix = 3,         // Role: Mix
  kEnd              // Invalid value
}

CachePlacement

Specifies the physical storage location of the cache.

enum class CachePlacement : uint32_t {
  kHost = 0U,             // The cache is placed in the host memory.
  kDevice = 1U,           // The cache is placed in the device memory.
}

CacheDesc

Represents the cache description.

struct CacheDesc {
  CachePlacement placement = CachePlacement::kDevice;    // Memory type.
  uint32_t num_tensors = 0U;                             // Number of tensors in the cache.
  DataType data_type = DT_UNDEFINED;                     // Data type of the tensor in the cache.
  std::vector<int64_t> shape;                            // Shape of the tensor in the cache.
  uint8_t reserved[128];                                 // Reserved.
}

CacheIndex

Cache index.

struct CacheIndex {
  uint64_t cluster_id;        // ID of the cluster where the cache is located.
  int64_t cache_id;           // Cache ID.
  uint32_t batch_index;       // Index of the batch used for PullKvCache.
  uint8_t reserved[128];      // Reserved.
}

Cache

Represents the cache information that contains the addresses of a set of tensors.

struct Cache {
  int64_t cache_id = -1;                     // Cache ID.
  std::vector<uintptr_t> tensor_addrs;       // Addresses of tensors in the cache. In the single-process multi-card scenario, the addresses of multiple cards are arranged in sequence.
  CacheDesc cache_desc;                      // Cache description.
  uint8_t reserved[128];                     // Reserved.
}

ClusterInfo and IpInfo

Represents the cluster information.

struct ClusterInfo {
  uint64_t remote_cluster_id = 0U;     // Cluster ID of the peer LLM-DataDist.
  int32_t remote_role_type = 0;        // role_type of the peer LLM-DataDist. 0 indicates the full role, and 1 indicates the incremental role.
  std::vector<IpInfo> local_ip_infos;  // IP address of the local LLM-DataDist. For details, see the following IpInfo structure.
  std::vector<IpInfo> remote_ip_infos; // IP address of the peer LLM-DataDist. For details, see the following IpInfo structure.
  uint8_t reserved[128];               // Reserved.
}

struct IpInfo {
  AscendString ip;         // IP address.
  uint16_t port = 0U;      // Port number.
  uint8_t reserved[128];   // Reserved.
}

KvCacheExtParam

Represents the extended parameters passed when the KV cache pull or push APIs are called.

struct KvCacheExtParam {
  std::pair<int32_t, int32_t> src_layer_range = {-1, -1};  // Layer range at the source side during KV transmission.
  std::pair<int32_t, int32_t> dst_layer_range{-1, -1};  // Layer range at the destination side during KV transmission.
  uint8_t tensor_num_per_layer = 2U;                       // Number of tensors per layer during KV transmission.
  uint8_t reserved[127];                                   // Reserved.
}

RegisterCfg

Represents the configuration parameters passed when the RegisterKvCache API is called.

struct RegisterCfg {
  uint8_t reserved[128] = {0}; // Reserved.
};

Parent topic: LLM-DataDist