aclnnAttentionToFFN
Applicable Products
| Product | Supported |
|---|---|
| × | |
| √ | |
| × | |
| × | |
| × | |
| × |
Function
Sends token data from the attention node to the FFN node.
Function Prototype
Each operator is divided into a two-phase API. You must first call the "aclnnAttentionToFFNGetWorkspaceSize" API to obtain the input parameters and calculate the required workspace size based on the computation process, and then call the "aclnnAttentionToFFN" API to perform the computation.
aclnnStatus aclnnAttentionToFFNGetWorkspaceSize(
const aclTensor *x,
const aclTensor *sessionId,
const aclTensor *microBatchId,
const aclTensor *layerId,
const aclTensor *expertIds,
const aclTensor *expertRankTable,
const aclTensor *scalesOptional,
const aclTensor *activeMaskOptional,
const char *group,
int64_t worldSize,
const aclIntArray *ffnTokenInfoTableShape,
const aclIntArray *ffnTokenDataShape,
const aclIntArray *attnTokenInfoTableShape,
int64_t moeExpertNum,
int64_t quantMode,
int64_t syncFlag,
int64_t ffnStartRankId,
uint64_t *workspaceSize,
aclOpExecutor **executor)aclnnStatus aclnnAttentionToFFN(
void *workspace,
uint64_t workspaceSize,
aclOpExecutor *executor,
aclrtStream stream)aclnnAttentionToFFNGetWorkspaceSize
Parameters:
Parameter Input/Output Description Instruction Data Type Data Format Dimension (Shape) Non-contiguous Tensor x Input Token data sent by the current card. The shape is (X, Bs, H). FLOAT16, BFLOAT16 ND 3 √ sessionId Input Indicates the ID of the current attention worker node. Shape is (X, ), with a value range of [0, attentionWorkerNum). INT32 ND 1 √ microBatchId Input Represents the ID of the micro batch. Shape: (X, ). INT32 ND 1 √ layerId Input Indicates the ID of the current model layer. Shape: (X,). INT32 ND 1 √ expertIds Input Indicates the topK expert indices for each token in each micro batch group. Shape: (X, Bs, K). Value range: [0, moeExpertNum). INT32 ND 3 √ expertRankTable Input Mapping table from expert IDs in each micro batch group to the expert deployment on FFN cards (the caller must ensure the values are correct). Shape: (L, moeExpertNum + sharedExpertNum, M). INT32 ND 3 √ scalesOptional Input Quantization smoothing parameter for each expert. Must be a null pointer in non-quantization scenarios. In dynamic quantization scenarios, can be valid data or a null pointer, with shape (L, moeExpertNum + sharedExpertNum, H). FLOAT32 ND 3 √ activeMaskOptional Input Indicates whether a token participates in communication. A valid value or a null pointer can be passed. When a null pointer is passed, all tokens participate in communication by default. When a value is passed, the shape is (X, Bs), where true indicates that the token participates in communication, and true must be placed before false. BOOL ND 2 √ group Input Name of the communicator (expert parallelism). String length range: [1, 128). STRING - - - worldSize Input Indicates the communicator size. Value range: [2, 768]. INT64 - - - ffnTokenInfoTableShape Input A list that represents the shape size of the token information table on the FFN node. The length is 3, including the number of attention nodes, the size of microBatchSize, and the shape size of the sending status information corresponding to each token. INT32 - - - ffnTokenDataShape Input Specifies a list of shape sizes of the token data table on the FFN node. Length is 5, including the number of attention nodes, microBatchSize, batchSize, number of experts to be sent per token (including shared experts), and the length of a single token. INT32 - - - attnTokenInfoTableShape Input A list representing the shape size of the token information table on the attention node. The length is 3, including the microBatchSize, batchSize, and the number of experts each token needs to send (including shared experts). INT32 - - - moeExpertNum Input Number of MoE experts. Value range: (0, 1024]. INT64 - - - quantMode Input Indicates the quantization mode. Only 0 (non-quantization) and 2 (dynamic quantization) are supported. INT64 - - - syncFlag Input Specifies the synchronization mode of the FFN node. Only 0 (synchronous) and 1 (asynchronous) are supported. INT64 - - - ffnStartRankId Input Indicates the start ID of the FFN node. The value range is [0, worldSize). INT64 - - - workspaceSize Output Returns the workspace size to be applied for on the device side. - - - - - executor Output Returns the op executor that contains the operator computation process. - - - - - Return value
aclnnStatus: return code. For details, see aclnn Return Code.
The first-phase API performs input parameter validation and returns an error in the following scenarios:
Return Value Return Code Description ACLNN_ERR_PARAM_NULLPTR 161001 The required input and output parameter Tensor is a null pointer. ACLNN_ERR_PARAM_INVALID 161002 The data types of the input and output are not within the supported range. ACLNN_ERR_INNER_TILING_ERROR 561002 The shapes of the input and output are not within the supported range. The parameter values are not within the supported range.
aclnnAttentionToFFN
Parameters:
Parameter Input/Output Description workspace Input Specifies the workspace memory address applied for on the Device side. workspaceSize Input Specifies the workspace size applied for on the Device side, which is obtained by the first-phase API aclnnAttentionToFFNGetWorkspaceSize.executor Input Operator executor that contains the operator computation flow. stream Input Specifies the stream for task execution. Return value
aclnnStatus: return code. For details, see aclnn Return Code.
Constraints
Deterministic Constraints:
aclnnAttentionToFFNdefaults to a deterministic implementation.
Parameter Consistency Constraints:
- The
group,worldSize,moeExpertNum,ffnTokenInfoTableShape,ffnTokenDataShape,ffnStartRankIdparameters and theHCCL_BUFFSIZEvalue used during the API call must be consistent across all cards, must also be consistent across different layers in the network, and must be consistent with the corresponding parameters of the operators in the decoupling scenario series.
- The
Product-Specific Constraints:
Atlas A3 training products/Atlas A3 inference products : In this scenario, a single card contains dual DIE (short for "die"), so "this card" in the parameter descriptions refers to a single DIE.
Shape Variable Constraints:
Variable Definition and Value Range X Indicates the micro batch sequence size (number of token groups). Currently, only X = 1is supported.H (hidden size) Indicates the hidden size. The value range is [1024, 8192].Bs Indicates the batch sequence size (number of tokens finally output by this card). The value range is 0 < Bs ≤ 512.K Indicates the number of top-K experts selected. The value range is 0 < K ≤ 16, and0 < K ≤ moeExpertNum.L Indicates the number of model layers. Currently, only L = 1is supported.M Indicates the length of the last dimension of expertRankTable, specifically the length of the expert deployment information list with the largest number of experts deployed on FFN nodes. The value range is 1 < M ≤ FFNWorkerNum * 2 + 1.moeExpertNum Indicates the number of MoE experts. The value range is (0, 1024].sharedExpertNum Indicates the number of shared experts (a shared expert can be replicated and deployed on multiple ffnRank cards). The value range is [0, 4].moeExpertRankNum Indicates the number of FFN nodes where MoE experts are deployed. The value range is 0 < moeExpertRankNum < FFNWorkerNum.sharedExpertRankNum Indicates the number of FFN nodes where shared experts are deployed. The value range is 0 < sharedExpertRankNum ≤ FFNWorkerNum.FFNWorkerNum Indicates the number of FFN nodes. The value range is 0 < FFNWorkerNum < worldSize, andFFNWorkerNum = moeExpertRankNum + sharedExpertRankNummust be satisfied.AttentionWorkerNum Indicates the number of Attention nodes. The value range is 0 < AttentionWorkerNum < worldSize.Environment Variable Constraints:
- HCCL_BUFFSIZE: Before calling this API, check whether the value of the HCCL_BUFFSIZE environment variable is reasonable. This environment variable indicates the memory size occupied by a single communicator, in MB. It defaults to 200 MB when not configured.
Communicator Usage Constraints:
- No other operators are allowed in the communicator of the AttentionToFFN operator.
Example
File Preparation:
Create an AttentionToFFNDemo directory, and create the aclnnAttentionToFFNDemo.cpp and AttentionToFFN.sh files in AttentionToFFNDemo as instructed below. Modify the code by referring to the following examples.
Install the CANN package, and compile and run AttentionToFFNDemo as instructed below.
AttentionToFFN.sh build script
#!/bin/bash
cann_path="/path/to/cann_env" # Change the path of the CANN environment.
g++ "aclnnAttentionToFFNDemo.cpp" -o AttentionToFFNDemo -I"$cann_path/910/include/" -I"$cann_path/910/include/aclnnop/" \
-L="$cann_path/910/lib64/" -lascendcl -lnnopbase -lopapi_math -lop_common -lpthread -lhcclCompile and run:
# Source the CANN environment.
source /path/to/cann_env/910/bin/setenv.bash
# Compile aclnnAttentionToFFNDemo.cpp
bash AttentionToFFN.sh
./AttentionToFFNDemoThe following is sample code for reference only.
#include <thread>
#include <iostream>
#include <string>
#include <cstring>
#include <vector>
#include "acl/acl.h"
#include "hccl/hccl.h"
#include "aclnn/opdev/fp16_t.h"
#include "aclnnop/aclnn_attention_to_ffn.h"
#define CHECK_RET(cond, return_expr) \
do { \
if (!(cond)) { \
return_expr; \
} \
} while (0)
#define LOG_PRINT(message, ...) \
do { \
printf(message, ##__VA_ARGS__); \
} while(0)
struct Args {
uint32_t rankId;
HcclComm hcclComm;
aclrtStream attentionToFFNStream;
aclrtContext context;
};
constexpr uint32_t WORLD_SIZE = 16;
constexpr uint32_t FFN_WORKER_NUM = 5;
constexpr uint32_t ATTENTION_WORKER_NUM = WORLD_SIZE - FFN_WORKER_NUM;
int64_t GetShapeSize(const std::vector<int64_t> &shape)
{
int64_t shape_size = 1;
for (auto i : shape) {
shape_size *= i;
}
return shape_size;
}
template<typename T>
int CreateAclTensor(const std::vector<T> &hostData, const std::vector<int64_t> &shape, void **deviceAddr,
aclDataType dataType, aclTensor **tensor)
{
auto size = GetShapeSize(shape) * sizeof(T);
auto ret = aclrtMalloc(deviceAddr, size, ACL_MEM_MALLOC_HUGE_FIRST);
CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("[ERROR] aclrtMalloc failed. ret: %d\n", ret); return ret);
ret = aclrtMemcpy(*deviceAddr, size, hostData.data(), size, ACL_MEMCPY_HOST_TO_DEVICE);
CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("[ERROR] aclrtMemcpy failed. ret: %d\n", ret); return ret);
std::vector<int64_t> strides(shape.size(), 1);
for (int64_t i = shape.size() - 2; i >= 0; i--) {
strides[i] = shape[i +1] * strides[i + 1];
}
*tensor = aclCreateTensor(shape.data(), shape.size(), dataType, strides.data(), 0,
aclFormat::ACL_FORMAT_ND, shape.data(), shape.size(), *deviceAddr);
return 0;
}
int LaunchOneProcessAttentionToFFN(Args &args)
{
int ret = aclrtSetCurrentContext(args.context);
CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("[ERROR] aclrtSetCurrentContext failed, ret %d\n", ret); return ret);
char hcomName[128] = {0};
ret = HcclGetCommName(args.hcclComm, hcomName);
CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("[ERROR] HcclGetCommName failed, ret %d\n", ret); return -1);
LOG_PRINT("[INFO] rank = %d, hcomName = %s, attentionToFFNStream = %p, context = %p\n",
args.rankId, hcomName, args.attentionToFFNStream, args.context);
int64_t X = 1;
int64_t L = 1;
int64_t Bs = 8;
int64_t H = 7168;
int64_t K = 4;
int64_t sharedExpertNum = 1;
int64_t sharedExpertRankNum = 1;
int64_t moeExpertNum = 8;
int64_t expertNumPerToken = K + sharedExpertNum;
int64_t M = 2 *(FFN_WORKER_NUM - sharedExpertNum * sharedExpertRankNum) + 1;
int64_t quantMode = 0;
int64_t syncFlag = 0;
int64_t ffnStartRankId = 0;
int64_t ffnTokenInfoTableShapeData[] = {ATTENTION_WORKER_NUM , X, 2 + Bs * expertNumPerToken};
int64_t ffnTokenDataShapeData[] = {ATTENTION_WORKER_NUM, X, Bs, expertNumPerToken, H};
int64_t attnTokenInfoTableShapeData[] = {X, Bs, expertNumPerToken};
/*Construct the input and output variables on the device side based on the current scenario.*/
// Declare the input and output variables on the device side.
void *xDeviceAddr = nullptr;
void *sessionIdDeviceAddr = nullptr;
void *microBatchIdDeviceAddr = nullptr;
void *layerIdDeviceAddr = nullptr;
void *expertIdsDeviceAddr = nullptr;
void *expertRankTableDeviceAddr = nullptr;
void *scalesDeviceAddr = nullptr;
aclTensor *x = nullptr;
aclTensor *sessionId = nullptr;
aclTensor *microBatchId = nullptr;
aclTensor *layerId = nullptr;
aclTensor *expertIds = nullptr;
aclTensor *expertRankTable = nullptr;
aclTensor *scales = nullptr;
aclIntArray *ffnTokenInfoTableShape = aclCreateIntArray(ffnTokenInfoTableShapeData, 3);
aclIntArray *ffnTokenDataShape = aclCreateIntArray(ffnTokenDataShapeData, 5);
aclIntArray *attnTokenInfoTableShape = aclCreateIntArray(attnTokenInfoTableShapeData, 3);
// Define the dimensions of each variable in the current scenario.
std::vector<int64_t> xShape{X, Bs, H};
std::vector<int64_t> sessionIdShape{X};
std::vector<int64_t> microBatchIdShape{X};
std::vector<int64_t> layerIdShape{X};
std::vector<int64_t> expertIdsShape{X, Bs, K};
std::vector<int64_t> expertRankTableShape{L, moeExpertNum + sharedExpertNum, M};
std::vector<int64_t> scalesShape{L, moeExpertNum + sharedExpertNum, H};
int64_t xShapeSize = GetShapeSize(xShape);
int64_t sessionIdShapeSize = GetShapeSize(sessionIdShape);
int64_t microBatchIdShapeSize = GetShapeSize(microBatchIdShape);
int64_t layerIdShapeSize = GetShapeSize(layerIdShape);
int64_t expertIdsShapeSize = GetShapeSize(expertIdsShape);
int64_t expertRankTableShapeSize = GetShapeSize(expertRankTableShape);
int64_t scalesShapeSize = GetShapeSize(scalesShape);
// Construct host-side variables.
std::vector<int16_t> xHostData(xShapeSize, 1);
std::vector<int16_t> sessionIdHostData(sessionIdShapeSize, args.rankId - FFN_WORKER_NUM);
std::vector<int16_t> microBatchIdHostData(microBatchIdShapeSize, 0);
std::vector<int16_t> layerIdHostData(layerIdShapeSize, 0);
std::vector<int32_t> expertIdsHostData;
for (int32_t micro_batch_id = 0; micro_batch_id < expertIdsShape[0]; micro_batch_id++) {
for (int32_t token_id = 0; token_id < expertIdsShape[1]; token_id++) {
for (int32_t k_id = 0; k_id < expertIdsShape[2]; k_id++) {
expertIdsHostData.push_back(k_id);
}
}
}
std::vector<int32_t> expertRankTableHostData = {4, 2, 4, 3, 7, 1, 3, 2, 5, 2, 2, 5, 1, 2, 0, 0, 0, 0,
3, 2, 5, 0, 0, 3, 7, 0, 0, 4, 1, 3, 0, 1, 2, 4, 3, 7,
4, 0, 0, 3, 6, 1, 3, 2, 5, 3, 3, 7, 2, 4, 1, 2, 0, 0,
2, 2, 5, 0, 0, 0, 0, 0, 0, 3, 3, 6, 2, 5, 3, 7, 0, 0,
1, 4, 8, 0, 0, 0, 0, 0, 0};
std::vector<float> scalesHostData(scalesShapeSize, 0.1);
ret = CreateAclTensor(xHostData, xShape, &xDeviceAddr, aclDataType::ACL_BF16, &x);
CHECK_RET(ret == ACL_SUCCESS, return ret);
ret = CreateAclTensor(sessionIdHostData, sessionIdShape, &sessionIdDeviceAddr, aclDataType::ACL_INT32, &sessionId);
CHECK_RET(ret == ACL_SUCCESS, return ret);
ret = CreateAclTensor(microBatchIdHostData, microBatchIdShape, µBatchIdDeviceAddr, aclDataType::ACL_INT32, µBatchId);
CHECK_RET(ret == ACL_SUCCESS, return ret);
ret = CreateAclTensor(layerIdHostData, layerIdShape, &layerIdDeviceAddr, aclDataType::ACL_INT32, &layerId);
CHECK_RET(ret == ACL_SUCCESS, return ret);
ret = CreateAclTensor(expertIdsHostData, expertIdsShape, &expertIdsDeviceAddr, aclDataType::ACL_INT32, &expertIds);
CHECK_RET(ret == ACL_SUCCESS, return ret);
ret = CreateAclTensor(expertRankTableHostData, expertRankTableShape, &expertRankTableDeviceAddr, aclDataType::ACL_INT32, &expertRankTable);
CHECK_RET(ret == ACL_SUCCESS, return ret);
ret = CreateAclTensor(scalesHostData, scalesShape, &scalesDeviceAddr, aclDataType::ACL_FLOAT, &scales);
CHECK_RET(ret == ACL_SUCCESS, return ret);
uint64_t attentionToFFNWorkspaceSize = 0;
aclOpExecutor *attentionToFFNExecutor = nullptr;
void *attentionToFFNWorkspaceAddr = nullptr;
/**************************************** Call AttentionToFFN ********************************************/
// Call the first-phase API.
ret = aclnnAttentionToFFNGetWorkspaceSize(x, sessionId, microBatchId, layerId, expertIds, expertRankTable, (quantMode > 0 ? scales : nullptr),
nullptr, hcomName, WORLD_SIZE, ffnTokenInfoTableShape, ffnTokenDataShape, attnTokenInfoTableShape,
moeExpertNum, quantMode, syncFlag, ffnStartRankId, &attentionToFFNWorkspaceSize, &attentionToFFNExecutor);
CHECK_RET(ret == ACL_SUCCESS,
LOG_PRINT("[ERROR] aclnnAttentionToFFNGetWorkspaceSize failed. ret = %d \n", ret); return ret);
// Apply for device memory based on the workspaceSize calculated by the first-phase API.
if (attentionToFFNWorkspaceSize > 0) {
ret = aclrtMalloc(&attentionToFFNWorkspaceAddr, attentionToFFNWorkspaceSize, ACL_MEM_MALLOC_HUGE_FIRST);
CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("[ERROR] aclrtMalloc workspace failed. ret = %d \n", ret); return ret);
}
if (args.rankId < FFN_WORKER_NUM) { // FFN Worker
// Wait for the Attention Worker task to finish.
LOG_PRINT("[INFO] device_%d is FFN worker, skipping aclnnAttentionToFFN execute.\n", args.rankId);
std::this_thread::sleep_for(std::chrono::seconds(30));
} else { // Attention Worker
// Call the second-phase API.
ret = aclnnAttentionToFFN(attentionToFFNWorkspaceAddr, attentionToFFNWorkspaceSize,
attentionToFFNExecutor, args.attentionToFFNStream);
// (Boilerplate) Synchronously wait for task execution to complete.
ret = aclrtSynchronizeStreamWithTimeout(args.attentionToFFNStream, 10000);
CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("[ERROR] aclnnAttentionToFFN failed. ret = %d \n", ret); \
return ret);
LOG_PRINT("[INFO] device_%d aclnnAttentionToFFN execute successfully.\n", args.rankId);
}
// Release device resources.
if (attentionToFFNWorkspaceSize > 0) {
aclrtFree(attentionToFFNWorkspaceAddr);
}
if (x != nullptr) {
aclDestroyTensor(x);
}
if (sessionId != nullptr) {
aclDestroyTensor(sessionId);
}
if (microBatchId != nullptr) {
aclDestroyTensor(microBatchId);
}
if (layerId != nullptr) {
aclDestroyTensor(layerId);
}
if (expertIds != nullptr) {
aclDestroyTensor(expertIds);
}
if (expertRankTable != nullptr) {
aclDestroyTensor(expertRankTable);
}
if (scales != nullptr) {
aclDestroyTensor(scales);
}
if (ffnTokenInfoTableShape != nullptr) {
aclDestroyIntArray(ffnTokenInfoTableShape);
}
if (ffnTokenDataShape != nullptr) {
aclDestroyIntArray(ffnTokenDataShape);
}
if (attnTokenInfoTableShape != nullptr) {
aclDestroyIntArray(attnTokenInfoTableShape);
}
if (xDeviceAddr != nullptr) {
aclrtFree(xDeviceAddr);
}
if (sessionIdDeviceAddr != nullptr) {
aclrtFree(sessionIdDeviceAddr);
}
if (microBatchIdDeviceAddr != nullptr) {
aclrtFree(microBatchIdDeviceAddr);
}
if (layerIdDeviceAddr != nullptr) {
aclrtFree(layerIdDeviceAddr);
}
if (expertIdsDeviceAddr != nullptr) {
aclrtFree(expertIdsDeviceAddr);
}
if (expertRankTableDeviceAddr != nullptr) {
aclrtFree(expertRankTableDeviceAddr);
}
if (scalesDeviceAddr != nullptr) {
aclrtFree(scalesDeviceAddr);
}
HcclCommDestroy(args.hcclComm);
aclrtDestroyStream(args.attentionToFFNStream);
aclrtDestroyContext(args.context);
aclrtResetDevice(args.rankId);
return 0;
}
int main(int argc, char *argv[])
{
// This sample is implemented based on Atlas A3 and must run on Atlas A3.
int ret = aclInit(nullptr);
CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("[ERROR] aclrtInit failed, ret = %d\n", ret); return ret);
aclrtStream attentionToFFNStream[WORLD_SIZE];
aclrtContext context[WORLD_SIZE];
for (uint32_t rankId = 0; rankId < WORLD_SIZE; rankId++) {
ret = aclrtSetDevice(rankId);
CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("[ERROR] aclrtSetDevice failed, ret = %d\n", ret); return ret);
ret = aclrtCreateContext(&context[rankId], rankId);
CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("[ERROR] aclrtCreateContext failed, ret = %d\n", ret); return ret);
ret = aclrtCreateStream(&attentionToFFNStream[rankId]);
CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("[ERROR] aclrtCreateStream failed, ret = %d\n", ret); return ret);
}
int32_t devices[WORLD_SIZE];
for (int32_t id = 0; id < WORLD_SIZE; id++) {
devices[id] = id;
}
HcclComm comms[WORLD_SIZE];
ret = HcclCommInitAll(WORLD_SIZE, devices, comms);
CHECK_RET(ret == ACL_SUCCESS,
LOG_PRINT("[ERROR] HcclCommInitAll failed, ret %d\n", ret); return ret);
Args args[WORLD_SIZE];
std::vector<std::unique_ptr<std::thread>> threads(WORLD_SIZE);
for (uint32_t rankId = 0; rankId < WORLD_SIZE; rankId++) {
args[rankId].rankId = rankId;
args[rankId].hcclComm = comms[rankId];
args[rankId].attentionToFFNStream = attentionToFFNStream[rankId];
args[rankId].context = context[rankId];
threads[rankId].reset(new(std::nothrow) std::thread(&LaunchOneProcessAttentionToFFN, std::ref(args[rankId])));
}
for(uint32_t rankId = 0; rankId < WORLD_SIZE; rankId++) {
threads[rankId]->join();
}
aclFinalize();
LOG_PRINT("[INFO] aclFinalize success\n");
return 0;
}