aclnnFFNToAttention

Applicable Products

Function

Sends data from the FFN node to the Attention node.

Function Prototype

Each operator is divided into a two-phase API. The "aclnnFFNToAttentionGetWorkspaceSize" API must be called first to obtain the workspace size required for computation and the executor that encapsulates the operator computation flow. Then, the "aclnnFFNToAttention" API is called to perform the computation.

aclnnStatus aclnnFFNToAttentionGetWorkspaceSize(
    const aclTensor   *x,
    const aclTensor   *sessionIds,
    const aclTensor   *microBatchIds,
    const aclTensor   *tokenIds,
    const aclTensor   *expertOffsets,
    const aclTensor   *actualTokenNum,
    const aclTensor   *attnRankTableOptional,
    const char        *group,
    int64_t            worldSize,
    const aclIntArray *tokenInfoTableShape,
    const aclIntArray *tokenDataShape,
    uint64_t          *workspaceSize,
    aclOpExecutor    **executor)

aclnnStatus aclnnFFNToAttention(
    void           *workspace,
    uint64_t        workspaceSize,
    aclOpExecutor  *executor,
    aclrtStream     stream)

aclnnFFNToAttentionGetWorkspaceSize

• Parameters:

  Parameter	Input/Output	Description	Instruction	Data Type	Data Format	Dimension (Shape)	Non-contiguous Tensor
  x	Input	Token data sent by the current card.	Shape: (Y, H).	FLOAT16, BFLOAT16	ND	2	√
  sessionIds	Input	Index of the Attention Worker node for each token.	Shape is (Y, ), with a value range of [0, attnRankNum-1].	INT32	ND	1	√
  microBatchIds	Input	Micro-batch index of each token.	Shape: (Y, ), value range: [0, MicroBatchNum-1].	INT32	ND	1	√
  tokenIds	Input	Token index of each token in the micro batch.	Shape: (Y, ), value range: [0, Bs-1].	INT32	ND	1	√
  expertOffsets	Input	Index of each token in PerTokenExpertNum within tokenInfoTableShape.	Its shape is (Y, ), with a value range of [0, ExpertNumPerToken-1].	INT32	ND	1	√
  actualTokenNum	Input	Indicates the actual total token count sent by the current card, which is a 1D tensor.	The shape is (1, ).	INT64	ND	1	√
  attnRankTableOptional	Optional input	Maps the card ID corresponding to each Attention Worker.	Attention Workers must be deployed consecutively starting from card 0. If a null pointer is passed, the default policy is used: the ID of each card serves as the ID of the corresponding Attention Worker, with a value range of [0, attnRankNum-1].	INT32	ND	1	√
  group	Input	Name of the communicator (expert parallelism).	String length: [1, 128).	STRING	-	-	×
  worldSize	Input	Communicator size.	worldSize value range: [2, 768].	INT64	-	-	×
  tokenInfoTableShape	Input	Size of the token information list.	Contains the micro-batch size (MicroBatchNum), batch size (Bs), and the number of experts per token (ExpertNumPerToken).	INT32	-	-	×
  tokenDataShape	Input	Size of the token data list.	Contains the micro-batch size (MicroBatchNum), batch size (Bs), the number of experts per token (ExpertNumPerToken), and the token and scale length (HS).	INT32	-	-	×
  workspaceSize	Output	Size of the workspace to be requested on the device side.	-	-	-	-	×
  executor	Output	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	Error Code	Description
  ACLNN_ERR_PARAM_NULLPTR	161001	The required input or output parameter Tensor is a null pointer.
  ACLNN_ERR_PARAM_INVALID	161002	The data types of the inputs and outputs are not within the supported range.

aclnnFFNToAttention

• Parameters:

  Parameter	Input/Output	Description
  workspace	Input	Workspace memory address applied for on the Device side.
  workspaceSize	Input	Workspace size applied for on the Device side, obtained from the first-phase API aclnnFFNToAttentionGetWorkspaceSize.
  executor	Input	Operator executor that contains the operator computation flow.
  stream	Input	Specifies the stream for task execution.

• Return value

Returns an aclnnStatus status code. For details, see aclnn return code.

Constraints

• Deterministic Constraints:

  • aclnnFFNToAttention defaults to a deterministic implementation.

• Parameter Consistency Constraints:

  • The group, worldSize, tokenInfoTableShape, tokenDataShape parameters and the HCCL_BUFFSIZE value must be consistent across all cards.

• Product-Specific Constraints:

  • Atlas A3 training products/Atlas A3 inference products: In this scenario, a single card contains dual DIE (referred to as "die" or "bare die"), so "this card" in the parameter description refers to a single DIE.

• Shape Variable Constraints:

  Variable	Definition and Value Range
  Y	Indicates the maximum number of tokens to be distributed by this card.
  Bs	Indicates the number of tokens sent on each Attention node. Atlas A3 training products/Atlas A3 inference products: 0 < Bs ≤ 512 .
  H (hidden size)	Indicates the hidden size. Atlas A3 training products/Atlas A3 inference products: 1024 ≤ H ≤ 8192 .
  HS (hidden and scale size)	Indicates the hidden and scale size. Atlas A3 training products/Atlas A3 inference products: 1152 ≤ HS ≤ 8320 .
  MicroBatchNum	Indicates the microBatch size. Currently, only MicroBatchNum = 1 is supported.
  ExpertNumPerToken	Indicates the number of experts sent per token, ExpertNumPerToken = K + sharedExpertNum.
  K	Indicates selection of the top-K experts. The value range is 0 < K ≤ 16 .
  ffnRankNum	Indicates selection of ffnRankNum cards as FFnWorkers. The value range is 0 < ffnRankNum < worldSize .
  attnRankNum	Indicates selection of attnRankNum cards as AttnWorkers. The value range is 0 < attnRankNum < worldSize .
  sharedExpertNum	Indicates the number of shared experts (a shared expert can be replicated and deployed on multiple ffnRank cards). The current value range is [0, 4].

• Communicator Usage Constraints:

  • No other operators are allowed in the communicator of the FFNToAttention operator.

Example

• File Preparation:

  1. Create the FFNtoAttentionDemo directory, and under it create aclnnFFNtoAttentionDemo.cpp and FFNtoAttention.sh as instructed below, then modify them by referring to the following code.

  2. Install the CANN package, and compile and run FFNtoAttentionDemo as instructed below.

• FFNtoAttention.sh build script

  #!/bin/bash
  cann_path="/path/to/cann_env" # Change the path to the CANN package environment.
  g++ "aclnnFFNtoAttentionDemo.cpp" -o FFNtoAttentionDemo -I"$cann_path/910/include/" -I"$cann_path/910/include/aclnnop/" \
                      -L="$cann_path/910/lib64/" -lascendcl -lnnopbase -lopapi_math -lop_common -lpthread -lhccl

• Compilation and execution:

  # Source the CANN environment.
  source /path/to/cann_env/910/bin/setenv.bash
  
  # Compile aclnnFFNtoAttentionDemo.cpp.
  bash FFNtoAttention.sh
  
  ./FFNtoAttentionDemo

• The sample code below is for reference only.

  #include <thread>
  #include <iostream>
  #include <string>
  #include <vector>
  #include <unordered_set>
  #include "acl/acl.h"
  #include "hccl/hccl.h"
  #include "aclnnop/aclnn_ffn_to_attention.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 FFN2AttentionStream;
      aclrtContext context;
  };
  
  constexpr uint32_t WORLD_SIZE = 16;
  constexpr uint32_t ATTN_NUM = 8;
  constexpr uint32_t DEV_NUM = WORLD_SIZE;
  
  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 LaunchOneProcessFFN2Attention(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, FFN2AttentionStream = %p, \
                  context = %p\n", args.rankId, hcomName, args.FFN2AttentionStream, \
                  args.context);
  
      int64_t micro_batch_num = 1;
      int64_t Y = 8;
      int64_t H = 7168;
      int64_t K = 7;
      int64_t attention_worker_num = ATTN_NUM;
      int64_t sharedExpertNum = 1;
      int64_t expert_num_per_token = K + sharedExpertNum;
      int64_t Token_info_shape[] = {micro_batch_num, Y, expert_num_per_token};
      int64_t Token_data_shape[] = {micro_batch_num, Y, expert_num_per_token, H};
  
      void *xDeviceAddr = nullptr;
      void *sessionIdsDeviceAddr = nullptr;
      void *microBatchIdsDeviceAddr = nullptr;
      void *tokenIdsDeviceAddr = nullptr;
      void *expertOffsetsDeviceAddr = nullptr;
      void *actualTokenNumDeviceAddr = nullptr;
      void *attnRankTableDeviceAddr = nullptr;
  
      aclTensor *x = nullptr;
      aclTensor *sessionIds = nullptr;
      aclTensor *microBatchIds = nullptr;
      aclTensor *tokenIds = nullptr;
      aclTensor *expertOffsets = nullptr;
      aclTensor *actualTokenNum = nullptr;
      aclTensor *attnRankTable = nullptr;
      aclIntArray *tokenInfoTableShape = aclCreateIntArray(Token_info_shape, 3);
      aclIntArray *tokenDataShape = aclCreateIntArray(Token_data_shape, 4);
  
  
      //Define the dimensions of each variable in the current scenario.
      std::vector<int64_t> xShape{Y, H};
      std::vector<int64_t> sessionIdsShape{Y};
      std::vector<int64_t> microBatchIdsShape{Y};
      std::vector<int64_t> tokenIdsShape{Y};
      std::vector<int64_t> expertOffsetsShape{Y};
      std::vector<int64_t> actualTokenNumShape{1};
      std::vector<int64_t> attnRankTableShape{attention_worker_num};
  
  
      int64_t xShapeSize = GetShapeSize(xShape);
      int64_t sessionIdsShapeSize = GetShapeSize(sessionIdsShape);
      int64_t microBatchIdsShapeSize = GetShapeSize(microBatchIdsShape);
      int64_t tokenIdsShapeSize = GetShapeSize(tokenIdsShape);
      int64_t expertOffsetsShapeSize = GetShapeSize(expertOffsetsShape);
      int64_t actualTokenNumShapeSize = GetShapeSize(actualTokenNumShape);
      int64_t attnRankTableShapeSize = GetShapeSize(attnRankTableShape);
  
  
      std::vector<int16_t> xHostData(xShapeSize, 1);
      std::vector<int32_t> sessionIdsHostData(sessionIdsShapeSize, 0);
      std::vector<int32_t> microBatchIdsHostData(microBatchIdsShapeSize, 0);
      std::vector<int32_t> tokenIdsHostData(tokenIdsShapeSize, 0);
      std::vector<int32_t> expertOffsetsHostData(expertOffsetsShapeSize, 0);
      std::vector<int64_t> actualTokenNumHostData(actualTokenNumShapeSize, 8);
      std::vector<int32_t> attnRankTableHostData(attnRankTableShapeSize);
      for (int32_t i = 0; i < Y; i++) {
          sessionIdsHostData[i] = i % attention_worker_num;
          tokenIdsHostData[i] = i % Y;
          expertOffsetsHostData[i] = i % expert_num_per_token;
      }
      for (int32_t i = 0; i < attention_worker_num; i++) {
          attnRankTableHostData[i] = static_cast<int32_t>(i);
      }
  
  
      ret = CreateAclTensor(xHostData, xShape, &xDeviceAddr, aclDataType::ACL_BF16, &x);
      CHECK_RET(ret == ACL_SUCCESS, return ret);
      ret = CreateAclTensor(sessionIdsHostData, sessionIdsShape, &sessionIdsDeviceAddr, aclDataType::ACL_INT32, &sessionIds);
      CHECK_RET(ret == ACL_SUCCESS, return ret);
      ret = CreateAclTensor(microBatchIdsHostData, microBatchIdsShape, µBatchIdsDeviceAddr, aclDataType::ACL_INT32, µBatchIds);
      CHECK_RET(ret == ACL_SUCCESS, return ret);
      ret = CreateAclTensor(tokenIdsHostData, tokenIdsShape, &tokenIdsDeviceAddr, aclDataType::ACL_INT32, &tokenIds);
      CHECK_RET(ret == ACL_SUCCESS, return ret);
      ret = CreateAclTensor(expertOffsetsHostData, expertOffsetsShape, &expertOffsetsDeviceAddr, aclDataType::ACL_INT32, &expertOffsets);
      CHECK_RET(ret == ACL_SUCCESS, return ret);
      ret = CreateAclTensor(actualTokenNumHostData, actualTokenNumShape, &actualTokenNumDeviceAddr,  aclDataType::ACL_INT64, &actualTokenNum);
      CHECK_RET(ret == ACL_SUCCESS, return ret);
      ret = CreateAclTensor(attnRankTableHostData, attnRankTableShape, &attnRankTableDeviceAddr, aclDataType::ACL_INT32, &attnRankTable);
      CHECK_RET(ret == ACL_SUCCESS, return ret);
  
  
      uint64_t FFN2AttentionWorkspaceSize = 0;
      aclOpExecutor *FFN2AttentionExecutor = nullptr;
      void *FFN2AttentionWorkspaceAddr = nullptr;
  
      /**************************************** Call FFN2Attention ********************************************/
      // Call the first-phase API.
      ret = aclnnFFNToAttentionGetWorkspaceSize(x, sessionIds, microBatchIds, tokenIds,
                                                          expertOffsets, actualTokenNum, attnRankTable,
                                                          hcomName, WORLD_SIZE, tokenInfoTableShape, tokenDataShape,
                                                          &FFN2AttentionWorkspaceSize, &FFN2AttentionExecutor);
      CHECK_RET(ret == ACL_SUCCESS,
          LOG_PRINT("[ERROR] aclnnFFNToAttentionGetWorkspaceSize failed. ret = %d \n", ret); return ret);
      // Apply for device memory based on the workspaceSize calculated by the first-phase API.
      if (FFN2AttentionWorkspaceSize > 0) {
          ret = aclrtMalloc(&FFN2AttentionWorkspaceAddr, FFN2AttentionWorkspaceSize, ACL_MEM_MALLOC_HUGE_FIRST);
          CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("[ERROR] aclrtMalloc workspace failed. ret = %d \n", ret); return ret);
      }
      // Call the second-phase API.
      ret = aclnnFFNToAttention(FFN2AttentionWorkspaceAddr, FFN2AttentionWorkspaceSize, FFN2AttentionExecutor, args.FFN2AttentionStream);
      CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("[ERROR] aclnnFFNToAttention failed. ret = %d \n", ret);
          return ret);
      // (Boilerplate) Synchronously wait for task execution to complete.
      if (args.rankId >= ATTN_NUM) {
          ret = aclrtSynchronizeStreamWithTimeout(args.FFN2AttentionStream, 10000);
          CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("[ERROR] aclrtSynchronizeStreamWithTimeout failed. ret = %d \n", ret);
          return ret);
          LOG_PRINT("[INFO] device_%d FFNToAttention execute successfully.\n", args.rankId);
      } else {
          std::this_thread::sleep_for(std::chrono::seconds(10));
          LOG_PRINT("[INFO] device_%d is AttentionWorker, sleeping 10 seconds...\n", args.rankId);
      }
  
  
      // Release device resources.
      if (FFN2AttentionWorkspaceSize > 0) {
          aclrtFree(FFN2AttentionWorkspaceAddr);
      }
      if (x != nullptr) {
          aclDestroyTensor(x);
      }
      if (sessionIds != nullptr) {
          aclDestroyTensor(sessionIds);
      }
      if (microBatchIds != nullptr) {
          aclDestroyTensor(microBatchIds);
      }
      if (tokenIds != nullptr) {
          aclDestroyTensor(tokenIds);
      }
  
      if (expertOffsets != nullptr) {
          aclDestroyTensor(expertOffsets);
      }
      if (actualTokenNum != nullptr) {
          aclDestroyTensor(actualTokenNum);
      }
      if (attnRankTable != nullptr) {
          aclDestroyTensor(attnRankTable);
      }
      if (tokenInfoTableShape != nullptr) {
          aclDestroyIntArray(tokenInfoTableShape);
      }
      if (tokenDataShape != nullptr) {
          aclDestroyIntArray(tokenDataShape);
      }
  
  
      if (xDeviceAddr != nullptr) {
          aclrtFree(xDeviceAddr);
      }
      if (sessionIdsDeviceAddr != nullptr) {
          aclrtFree(sessionIdsDeviceAddr);
      }
      if (microBatchIdsDeviceAddr != nullptr) {
          aclrtFree(microBatchIdsDeviceAddr);
      }
      if (tokenIdsDeviceAddr != nullptr) {
          aclrtFree(tokenIdsDeviceAddr);
      }
      if (expertOffsetsDeviceAddr != nullptr) {
          aclrtFree(expertOffsetsDeviceAddr);
      }
      if (actualTokenNumDeviceAddr != nullptr) {
          aclrtFree(actualTokenNumDeviceAddr);
      }
      if (attnRankTableDeviceAddr != nullptr) {
          aclrtFree(attnRankTableDeviceAddr);
      }
  
      HcclCommDestroy(args.hcclComm);
      aclrtDestroyStream(args.FFN2AttentionStream);
      aclrtDestroyContext(args.context);
      aclrtResetDevice(args.rankId);
  
      return 0;
  }
  
  int main(int argc, char *argv[])
  {
      int ret = aclInit(nullptr);
      CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("[ERROR] aclInit failed, ret = %d\n", ret); return ret);
  
      aclrtStream FFN2AttentionStream[DEV_NUM];
      aclrtContext context[DEV_NUM];
      for (uint32_t rankId = 0; rankId < DEV_NUM; 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(&FFN2AttentionStream[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 deviceId = 0; deviceId < WORLD_SIZE; deviceId++) {
          devices[deviceId] = deviceId ;
      }
  
      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[DEV_NUM];
      std::vector<std::unique_ptr<std::thread>> threads(DEV_NUM);
      for (uint32_t rankId = 0; rankId < DEV_NUM; rankId++) {
          args[rankId].rankId = rankId;
          args[rankId].hcclComm = comms[rankId];
          args[rankId].FFN2AttentionStream = FFN2AttentionStream[rankId];
          args[rankId].context = context[rankId];
          threads[rankId].reset(new(std::nothrow) std::thread(&LaunchOneProcessFFN2Attention, std::ref(args[rankId])));
      }
  
      for(uint32_t rankId = 0; rankId < DEV_NUM; rankId++) {
          threads[rankId]->join();
      }
  
      aclFinalize();
      LOG_PRINT("[INFO] aclFinalize success\n");
  
      return 0;
  }