aclnnMlaPrologV2WeightNz

Note: This API will be deprecated in later versions. Use the latest API aclnnMlaPrologV3WeightNz instead.

Applicable Products

Product	Supported
Ascend 950PR/Ascend 950DT	×
Atlas A3 training products/Atlas A3 inference products	√
Atlas A2 training products/Atlas A2 inference products	√

Function

Description: In inference scenarios, this API performs the pre-processing computation for Multi-Head Latent Attention. The main computation process is divided into five paths:
- First, the input $x$ is multiplied by $W^{DQ}$ for downsampling and RmsNorm, and then split into two paths. The first path is multiplied by $W^{UQ}$ and $W^{UK}$ and undergoes two upsampling operations to obtain $q^N$ . The second path is multiplied by $W^{QR}$ and then undergoes rotary position encoding (ROPE) to obtain $q^R$ .
- The third path: the input $x$ is multiplied by $W^{DKV}$ for downsampling and RmsNorm, then passed into the Cache to obtain $k^C$ ;
- The fourth path: the input $x$ is multiplied by $W^{KR}$ , then undergoes rotary position encoding and is passed into another Cache to obtain $k^R$ ;
- The fifth path: the quantization parameter obtained after the output $q^N$ undergoes DynamicQuant.
- The weight parameters WeightDq, WeightUqQr, and WeightDkvKr must be passed in NZ format.
Formula:

RmsNorm formula
$\text{RmsNorm}(x) = \gamma \cdot \frac{x_i}{\text{RMS}(x)}$ $\text{RMS}(x) = \sqrt{\frac{1}{N} \sum_{i=1}^{N} x_i^2 + \epsilon}$
The computation formula of Query includes downsampling, RmsNorm, and two upsampling operations.
$c^Q = RmsNorm(x \cdot W^{DQ})$ $q^C = c^Q \cdot W^{UQ}$ $q^N = q^C \cdot W^{UK}$
Perform ROPE rotary positional encoding on Query.
$q^R = ROPE(c^Q \cdot W^{QR})$
Key computation formula, including downsampling and RmsNorm, with the computation result stored in cache.
$c^{KV} = RmsNorm(x \cdot W^{DKV})$ $k^C = Cache(c^{KV})$
Perform ROPE rotary positional encoding on Key and store the result in cache.
$k^R = Cache(ROPE(x \cdot W^{KR}))$
Dequant Scale Query Nope calculation formula:
$dequantScaleQNope = {RowMax(abs(q^{N})) / 127}$ $q^{N} = {round(q^{N} / dequantScaleQNope)}$

Function Prototype

Each operator uses a two-phase API. You must first call the "aclnnMlaPrologV2WeightNzGetWorkspaceSize" API to obtain the input parameters and calculate the required workspace size based on the process, and then call the "aclnnMlaPrologV2WeightNz" API to perform the computation.

aclnnStatus aclnnMlaPrologV2WeightNzGetWorkspaceSize(
  const aclTensor *tokenX,
  const aclTensor *weightDq,
  const aclTensor *weightUqQr,
  const aclTensor *weightUk,
  const aclTensor *weightDkvKr,
  const aclTensor *rmsnormGammaCq,
  const aclTensor *rmsnormGammaCkv,
  const aclTensor *ropeSin,
  const aclTensor *ropeCos,
  const aclTensor *cacheIndex,
  aclTensor       *kvCacheRef,
  aclTensor       *krCacheRef,
  const aclTensor *dequantScaleXOptional,
  const aclTensor *dequantScaleWDqOptional,
  const aclTensor *dequantScaleWUqQrOptional,
  const aclTensor *dequantScaleWDkvKrOptional,
  const aclTensor *quantScaleCkvOptional,
  const aclTensor *quantScaleCkrOptional,
  const aclTensor *smoothScalesCqOptional,
  double           rmsnormEpsilonCq,
  double           rmsnormEpsilonCkv,
  char            *cacheModeOptional,
  const aclTensor *queryOut,
  const aclTensor *queryRopeOut,
  const aclTensor *dequantScaleQNopeOutOptional,
  uint64_t        *workspaceSize,
  aclOpExecutor  **executor)

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

aclnnMlaPrologV2WeightNzGetWorkspaceSize

Parameters

Parameter	Input/Output	Description	Instruction	Data Type	Data Format	Dimension (Shape)	Non-contiguous Tensor
tokenX	Input	Input tensor used to compute Query and Key in the formula.	- Supports empty tensor with B=0, S=0, T=0.	BFLOAT16, INT8	ND	- Concatenated axis: (T,He) - Non-concatenated axis: (B,S,He)	×
weightDq	Input	Downsampling weight matrix $W^{DQ}$ used to compute Query in the formula.	- Null tensor not supported.	BFLOAT16, INT8	FRACTAL_NZ	(He,Hcq)	×
weightUqQr	Input	Upsampling weight matrix $W^{UQ}$ and positional encoding weight matrix $W^{QR}$ used to compute Query in the formula.	- Null tensor not supported. When dtype is INT8 (quantization scenario): 1. Must be per-tensor quantization input. 2. dequantScaleWUqQrOptional must be passed for non-quantized output. 3. dequantScaleWUqQrOptional, quantScaleCkvOptional, and quantScaleCkrOptional must be passed for quantized output. 4. smoothScalesCqOptional is optional. When dtype is BFLOAT16 (non-quantized scenario): 1. dequantScaleWUqQrOptional, quantScaleCkvOptional, quantScaleCkrOptional, and smoothScalesCqOptional must be null pointers.	BFLOAT16, INT8	FRACTAL_NZ	(Hcq,N*(D+Dr))	×
weightUk	Input	Upsampling weight $W^{UK}$ used to compute Key.	- Null tensor not supported.	BFLOAT16	ND	(N,D,Hckv)	×
weightDkvKr	Input	Downsampling weight matrix $W^{DKV}$ and positional encoding weight matrix $W^{KR}$ used to compute Key.	- Null tensor not supported.	BFLOAT16, INT8	FRACTAL_NZ	(He,Hckv+Dr)	×
rmsnormGammaCq	Input	The $\gamma$ parameter in the RmsNorm formula for computing $c^Q$ .	- Null tensor not supported.	BFLOAT16	ND	(Hcq)	×
rmsnormGammaCkv	Input	The $\gamma$ parameter in the RmsNorm formula for computing $c^{KV}$ .	- Null tensor not supported.	BFLOAT16	ND	(Hckv)	×
ropeSin	Input	Sine parameter matrix used to compute rotary position encoding.	- Supports empty tensor with B=0, S=0, T=0.	BFLOAT16	ND	- Concatenated axis: (T,Dr) - Non-concatenated axis: (B,S,Dr)	×
ropeCos	Input	Cosine parameter matrix used to compute rotary position encoding.	- Supports empty tensor with B=0, S=0, T=0.	BFLOAT16	ND	- Concatenated axis: (T,Dr) - Non-concatenated axis: (B,S,Dr)	×
cacheIndex	Input	Index used to store kvCache and krCache.	- Supports empty tensor with B=0, S=0, T=0. The value range must be within [0, BlockNum*BlockSize).	INT64	ND	- Concatenated axis: (T) - Non-concatenated axis: (B,S)	×
kvCacheRef	Input	aclTensor used for cache indexing. The computation result is updated in-place (corresponding to $k^C$ in the formula).	- Supports empty tensor with B=0, Skv=0. Nkv is associated with N, where N is a hyperparameter, so Nkv cannot be 0.	BFLOAT16, INT8	ND	(BlockNum,BlockSize,Nkv,Hckv)	×
krCacheRef	Input	Cache used for key positional encoding. The computation result is updated in-place (corresponding to $k^R$ in the formula).	- Supports empty tensor with B=0, Skv=0. Nkv is associated with N, where N is a hyperparameter, so Nkv cannot be 0.	BFLOAT16, INT8	ND	(BlockNum,BlockSize,Nkv,Dr)	×
dequantScaleXOptional	Input	Dequantization parameter of tokenX.	- Data format supports ND.	FLOAT	ND	- Concatenated axis: (T) - Non-concatenated axis: (B*S,1)	×
dequantScaleWDqOptional	Input	Dequantization parameter of weightDq.	- Data format supports ND.	FLOAT	ND	(1,Hcq)	×
dequantScaleWUqQrOptional	Input	Per-channel parameter used for dequantization after the MatmulQcQr matrix multiplication.	- Non-null tensor supported (required only for INT8 dtype scenario).	FLOAT	ND	(1,N*(D+Dr))	×
dequantScaleWDkvKrOptional	Input	Dequantization parameter of weightDkvKr.	- Data format supports ND.	FLOAT	ND	(1, Hckv+Dr)	×
quantScaleCkvOptional	Input	Parameter used to quantize the kvCache output data.	- Non-null tensor supported (required only for INT8 dtype quantized output scenario).	FLOAT	ND	(1,Hckv)	×
quantScaleCkrOptional	Input	Parameter used to quantize the krCache output data.	- Non-null tensor supported (required only for INT8 dtype quantized output scenario).	FLOAT	ND	(1,Dr)	×
smoothScalesCqOptional	Input	Parameter used for dynamic quantization of the RmsNormCq output.	- Non-null tensor supported (optional for INT8 dtype scenario).	FLOAT	ND	(1,Hcq)	×
rmsnormEpsilonCq	Input	The $\epsilon$ parameter in the RmsNorm formula for computing $c^Q$ .	- 1e-05 is recommended when not explicitly specified by the user. - Only the double type is supported.	DOUBLE	-	-	-
rmsnormEpsilonCkv	Input	The $\epsilon$ parameter in the RmsNorm formula for computing $c^{KV}$ .	- 1e-05 is recommended when not explicitly specified by the user. - Only the double type is supported.	DOUBLE	-	-	-
cacheModeOptional	Input	Indicates the mode of kvCache.	- "PA_BSND" is recommended when not explicitly specified by the user. - Only the char* type is supported. - Valid values are "PA_BSND" and "PA_NZ".	CHAR*	-	-	-
queryOut	Output	Output tensor of Query in the formula (corresponding to $q^N$ ).	-	BFLOAT16, INT8	ND	- Concatenated axis: (T,N,Hckv) - Non-concatenated axis: (B,S,N,Hckv)	×
queryRopeOut	Output	Output tensor of Query positional encoding in the formula (corresponding to $q^R$ ).	-	BFLOAT16	ND	- Concatenated axis: (T,N,Dr) - Non-concatenated axis: (B,S,N,Dr)	×
dequantScaleQNopeOutOptional	Output	Dequantization parameter of the Query output.	-	FLOAT	ND	- Concatenated axis: (T,1) - Non-concatenated axis: (B*S,1)	×
workspaceSize	Output	Returns the workspace size to be allocated on the Device side.	- Used only for outputting the result; no input configuration is required. - The data type is uint64_t*.	-	-	-	-
executor	Output	Returns the op executor, which contains the operator computation flow.	- Used only for outputting the result; no input configuration is required. - The data type is aclOpExecutor**.	-	-	-	-

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	A null pointer exists in a parameter that must be passed (such as an input/output parameter that the API core depends on).
ACLNN_ERR_PARAM_INVALID	161002	The shape (dimension/size) or dtype (data type) of an input parameter is not within the supported range of the API.
ACLNN_ERR_RUNTIME_ERROR	361001	An exception occurs when the API memory calls the NPU Runtime API (such as the Runtime service not started, memory allocation failure, etc.).
ACLNN_ERR_INNER_TILING_ERROR	561002	A tiling exception occurs, where the dtype type or shape of an input parameter is incorrect.

aclnnMlaPrologV2WeightNz

Parameters

Name	Type	Meaning
workspace	void*	Device-side allocated workspace memory address.
workspaceSize	uint64_t	Device-side allocated workspace size, obtained from the first-phase API aclnnMlaPrologV2WeightNzGetWorkspaceSize.
executor	aclOpExecutor*	Op executor that contains the operator computation flow.
stream	aclrtStream	Stream on which the task is executed.

Return Value

aclnnStatus: return code. For details, see aclnn Return Code.

Constraints

Deterministic computation:
- aclnnMlaPrologV2WeightNz defaults to a deterministic implementation.

Currently, A5 does not fully support the input parameter generalization range and API support scenarios of MlaProlog on A2 and A3. The details are as follows:

Shape format field description

Field Name	Full Name/Meaning	Value Rules and Description
B	Batch (input sample batch size)	Value range: 0 to 65536.
S	Seq-Length (input sample sequence length)	A2, A3 value range: unlimited. A5 value range: 0, 1.
He	Hidden-Size (hidden layer size)	A2, A3 fixed values: 1024, 2048, 3072, 4096, 5120, 6144, 7168, 7680, 8192. A5 fixed value: 7168.
Hcq	q low-rank matrix dimension	Fixed value: 1536.
N	Head-Num (number of heads)	Value range: 1, 2, 4, 8, 16, 32, 64, 128.
Hckv	kv low-rank matrix dimension	Fixed value: 512.
D	qk dimension excluding positional encoding	Fixed value: 128.
Dr	qk positional encoding dimension	Fixed value: 64.
Nkv	Number of kv heads	Fixed value: 1.
BlockNum	Number of blocks in PagedAttention scenario	Value is the result of `B*Skv/BlockSize` rounded up (Skv denotes the kv sequence length, 0 allowed).
BlockSize	Block size in PagedAttention scenario	Value range: 16 to 1024, and a multiple of 16.
T	Size after BS concatenated axis	A2, A3 value range: 0 to 1048576. Note: If BS concatenated axis is used, tokenX, ropeSin, and ropeCos are 2-dimensional, cacheIndex is 1-dimensional, and queryOut and queryRopeOut are 3-dimensional. A5 does not support BS concatenated axis yet.

When weight_dq, weight_uq_qr, and weight_dkv_kr are not transposed, the representation of each dimension is: (k, n).
Shape constraints:
- If the dimension of tokenX uses BS concatenated axis, i.e., (T, He)
  - The shape of ropeSin and ropeCos is (T, Dr)
  - The shape of cacheIndex is (T).
  - The shape of dequantScaleXOptional is (T, 1).
  - The shape of queryOut is (T, N, Hckv).
  - The shape of queryRopeOut is (T, N, Dr).
  - In the int8 full quantization scenario, the shape of dequantScaleQNopeOutOptional is (T, N, 1); in other scenarios, it is (1).
- If the dimension of tokenX does not use BS concatenated axis, i.e., (B, S, He)
  - The shape of ropeSin and ropeCos is (B, S, Dr)
  - The shape of cacheIndex is (B, S)
  - The shape of dequantScaleXOptional is (B*S, 1)
  - The shape of queryOut is (B, S, N, Hckv)
  - The shape of queryRopeOut is (B, S, N, Dr).
  - In the int8 full quantization scenario, the shape of dequantScaleQNopeOutOptional is (B*S, N, 1); in other scenarios, it is (1).
- One or more of B, S, T, and Skv may be 0, meaning that input parameters whose shapes are related to B, S, T, or Skv may be null tensors. The remaining input parameters do not support null tensors.
  - If B, S, and T are 0, queryOut and queryRopeOut output null tensors, and kvCacheRef and krCacheRef are not updated.
  - If Skv is 0, queryOut, queryRopeOut, and dequantScaleQNopeOutOptional are computed normally, while kvCacheRef and krCacheRef are not updated, i.e., they output null tensors.

The aclnnMlaPrologV2WeightNz API supports the following scenarios: A2 and A3 support all the following scenarios, while A5 currently only supports non-quantized scenarios.

Scenario		Meaning
Non-quantized		Input parameters: all input parameters are non-quantized data. Output parameters: all output parameters are non-quantized data.
Partial quantization	kv_cache non-quantized	Input parameters: weightUqQr passes per-token quantized data, and the remaining input parameters are all non-quantized data. Output parameters: all output parameters return non-quantized data.
Partial quantization	kv_cache quantized	Input parameters: weightUqQr passes per-token quantized data, kvCacheRef and krCacheRef pass per-channel quantized data, and the remaining input parameters are all non-quantized data. Output parameters: kvCacheRef and krCacheRef return per-channel quantized data, and the remaining output parameters return non-quantized data.
int8 full quantization	kv_cache non-quantized	Input parameters: tokenX passes pertoken quantized data, weightDq, weightUqQr, and weightDkvKr pass perchannel quantized data, and the remaining input parameters are all non-quantized data. Output parameters: all output parameters are non-quantized data.
int8 full quantization	kv_cache quantized	Input parameters: tokenX passes pertoken quantized data, weightDq, weightUqQr, and weightDkvKr pass perchannel quantized data, kvCacheRef passes pertensor quantized data, and the remaining input parameters are all non-quantized data. Output parameters: queryOut returns pertoken_head quantized data, kvCacheRef output returns pertensor quantized data, and the remaining output parameters are non-quantized data.

In different quantization scenarios, the dtype and shape combinations of parameters must meet the following conditions:

Parameter	Non-quantized Scenario		Partial Quantization Scenario				int8 full quantization scenario
	Non-quantized Scenario		kv_cache non-quantized		kv_cache quantized		kv_cache non-quantized		kv_cache quantized
	data type	shape	data type	shape	data type	Shape	Dtype	Shape	Dtype	Shape
tokenX	BFLOAT16	· (B,S,He) · (T, He)	BFLOAT16	· (B,S,He) · (T, He)	BFLOAT16	· (B,S,He) · (T, He)	INT8	· (B,S,He) · (T, He)	INT8	· (B,S,He) · (T, He)
weightDq	BFLOAT16	(He, Hcq)	BFLOAT16	(He, Hcq)	BFLOAT16	(He, Hcq)	INT8	(He, Hcq)	INT8	(He, Hcq)
weightUqQr	BFLOAT16	(Hcq, N*(D+Dr))	INT8	(Hcq, N*(D+Dr))	INT8	(Hcq, N*(D+Dr))	INT8	(Hcq, N*(D+Dr))	INT8	(Hcq, N*(D+Dr))
weightUk	BFLOAT16	(N, D, Hckv)	BFLOAT16	(N, D, Hckv)	BFLOAT16	(N, D, Hckv)	BFLOAT16	(N, D, Hckv)	BFLOAT16	(N, D, Hckv)
weightDkvKr	BFLOAT16	(He, Hckv+Dr)	BFLOAT16	(He, Hckv+Dr)	BFLOAT16	(He, Hckv+Dr)	INT8	(He, Hckv+Dr)	INT8	(He, Hckv+Dr)
rmsnormGammaCq	BFLOAT16	(Hcq)	BFLOAT16	(Hcq)	BFLOAT16	(Hcq)	BFLOAT16	(Hcq)	BFLOAT16	(Hcq)
rmsnormGammaCkv	BFLOAT16	(Hckv)	BFLOAT16	(Hckv)	BFLOAT16	(Hckv)	BFLOAT16	(Hckv)	BFLOAT16	(Hckv)
ropeSin	BFLOAT16	· (B,S,Dr) · (T, Dr )	BFLOAT16	· (B,S,Dr) · (T, Dr )	BFLOAT16	· (B,S,Dr) · (T, Dr )	BFLOAT16	· (B,S,Dr) · (T, Dr )	BFLOAT16	· (B,S,Dr) · (T, Dr )
ropeCos	BFLOAT16	· (B,S,Dr) · (T, Dr )	BFLOAT16	· (B,S,Dr) · (T, Dr )	BFLOAT16	· (B,S,Dr) · (T, Dr )	BFLOAT16	· (B,S,Dr) · (T, Dr )	BFLOAT16	· (B,S,Dr) · (T, Dr )
cacheIndex	INT64	· (B,S) · (T)	INT64	· (B,S) · (T)	INT64	· (B,S) · (T)	INT64	· (B,S) · (T)	INT64	· (B,S) · (T)
kvCacheRef	BFLOAT16	(BlockNum, BlockSize, Nkv, Hckv)	BFLOAT16	(BlockNum, BlockSize, Nkv, Hckv)	INT8	(BlockNum, BlockSize, Nkv, Hckv)	BFLOAT16	(BlockNum, BlockSize, Nkv, Hckv)	INT8	(BlockNum, BlockSize, Nkv, Hckv)
krCacheRef	BFLOAT16	(BlockNum, BlockSize, Nkv, Dr)	BFLOAT16	(BlockNum, BlockSize, Nkv, Dr)	INT8	(BlockNum, BlockSize, Nkv, Dr)	BFLOAT16	(BlockNum, BlockSize, Nkv, Dr)	BFLOAT16	(BlockNum, BlockSize, Nkv, Dr)
dequantScaleXOptional	No assignment required.	/	No assignment required	/	No assignment required	/	FLOAT	· (B*S, 1) · (T, 1)	FLOAT	· (B*S, 1) · (T, 1)
dequantScaleWDqOptional	No assignment required.	/	No assignment required.	/	No assignment required.	/	FLOAT	(1, Hcq)	FLOAT	(1, Hcq)
dequantScaleWUqQrOptional	No assignment required	/	FLOAT	(1, N*(D+Dr))	FLOAT	(1, N*(D+Dr))	FLOAT	(1, N*(D+Dr))	FLOAT	(1, N*(D+Dr))
dequantScaleWDkvKrOptional	No assignment required.	/	No assignment required.	/	No assignment required	/	FLOAT	(1, Hckv+Dr)	FLOAT	(1, Hckv+Dr)
quantScaleCkvOptional	No assignment required	/	No assignment required	/	FLOAT	(1, Hckv)	No assignment required.	/	FLOAT	(1, Hckv)
quantScaleCkrOptional	No assignment required	/	No assignment required	/	FLOAT	(1, Dr)	No assignment required	/	No assignment required.	/
smoothScalesCqOptional	No assignment required.	/	FLOAT	(1, Hcq)	FLOAT	(1, Hcq)	FLOAT	(1, Hcq)	FLOAT	(1, Hcq)
queryOut	BFLOAT16	· (B, S, N, Hckv) · (T, N, Hckv)	BFLOAT16	· (B, S, N, Hckv) · (T, N, Hckv)	BFLOAT16	· (B, S, N, Hckv) · (T, N, Hckv)	BFLOAT16	· (B, S, N, Hckv) · (T, N, Hckv)	INT8	· (B, S, N, Hckv) · (T, N, Hckv)
queryRopeOut	BFLOAT16	· (B, S, N, Dr) · (T, N, Dr)	BFLOAT16	· (B, S, N, Dr) · (T, N, Dr)	BFLOAT16	· (B, S, N, Dr) · (T, N, Dr)	BFLOAT16	· (B, S, N, Dr) · (T, N, Dr)	BFLOAT16	· (B, S, N, Dr) · (T, N, Dr)
dequantScaleQNopeOutOptional	No assignment required.	/	No assignment required.	/	No assignment required.	/	No assignment required.	/	FLOAT	· (B*S, N, 1) · (T, N, 1)

Example

The following provides A2 and A3 example code for reference only. For details about compilation and running, see Compile and Run Samples.

#include <iostream>
#include <vector>
#include <cstdint>
#include "acl/acl.h"
#include "aclnnop/aclnn_mla_prolog_v2_weight_nz.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)

int64_t GetShapeSize(const std::vector<int64_t>& shape) {
    int64_t shape_size = 1;
    for (auto i : shape) {
        shape_size *= i;
    }
    return shape_size;
}

int Init(int32_t deviceId, aclrtStream* stream) {
    // Boilerplate for resource initialization.
    auto ret = aclInit(nullptr);
    CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("aclInit failed. ERROR: %d\n", ret); return ret);
    ret = aclrtSetDevice(deviceId);
    CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("aclrtSetDevice failed. ERROR: %d\n", ret); return ret);
    ret = aclrtCreateStream(stream);
    CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("aclrtCreateStream failed. ERROR: %d\n", ret);  return ret);
    return 0;
}

template <typename T>
int CreateAclTensorND(const std::vector<T>& shape, void** deviceAddr, void** hostAddr,
                    aclDataType dataType, aclTensor** tensor) {
    auto size = GetShapeSize(shape) * sizeof(T);
    // Call aclrtMalloc to apply for device memory.
    auto ret = aclrtMalloc(deviceAddr, size, ACL_MEM_MALLOC_HUGE_FIRST);
    CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("aclrtMalloc failed. ERROR: %d\n", ret); return ret);
    // Call aclrtMalloc to apply for host memory.
    ret = aclrtMalloc(hostAddr, size,   ACL_MEM_MALLOC_HUGE_FIRST);
    CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("aclrtMalloc failed. ERROR: %d\n", ret); return ret);
    // Call the aclCreateTensor API to create an aclTensor.
    *tensor = aclCreateTensor(shape.data(), shape.size(), dataType, nullptr, 0, aclFormat::ACL_FORMAT_ND,
                              shape.data(), shape.size(), *deviceAddr);
    // Call aclrtMemcpy to copy host data to device memory.
    ret = aclrtMemcpy(*deviceAddr, size, *hostAddr, GetShapeSize(shape)*aclDataTypeSize(dataType), ACL_MEMCPY_HOST_TO_DEVICE);
    CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("aclrtMemcpy failed. ERROR: %d\n", ret); return ret);
    return 0;
}

template <typename T>
int CreateAclTensorNZ(const std::vector<T>& shape, void** deviceAddr, void** hostAddr,
                    aclDataType dataType, aclTensor** tensor) {
    auto size = GetShapeSize(shape) * sizeof(T);
    // Call aclrtMalloc to apply for device memory.
    auto ret = aclrtMalloc(deviceAddr, size, ACL_MEM_MALLOC_HUGE_FIRST);
    CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("aclrtMalloc failed. ERROR: %d\n", ret); return ret);
    // Call aclrtMalloc to apply for host-side memory.
    ret = aclrtMalloc(hostAddr, size, ACL_MEM_MALLOC_HUGE_FIRST);
    CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("aclrtMalloc failed. ERROR: %d\n", ret); return ret);
    // Call the aclCreateTensor API to create an aclTensor.
    *tensor = aclCreateTensor(shape.data(), shape.size(), dataType, nullptr, 0, aclFormat::ACL_FORMAT_FRACTAL_NZ,
                              shape.data(), shape.size(), *deviceAddr);
    // Call aclrtMemcpy to copy data from the host side to the device memory.
    ret = aclrtMemcpy(*deviceAddr, size, *hostAddr, GetShapeSize(shape)*aclDataTypeSize(dataType), ACL_MEMCPY_HOST_TO_DEVICE);
    CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("aclrtMemcpy  failed. ERROR: %d\n", ret); return ret);
    return 0;
}

int TransToNZShape(std::vector<int64_t> &shapeND, size_t typeSize) {
    int64_t h = shapeND[0];
    int64_t w = shapeND[1];
    int64_t h0 = 16;
    int64_t w0 = 32U / typeSize;
    int64_t h1 = h / h0;
    int64_t w1 = w / w0;
    shapeND[0] = w1;
    shapeND[1] = h1;
    shapeND.emplace_back(h0);
    shapeND.emplace_back(w0);
    return 0;
}

int main() {
    // 1. Boilerplate: device/stream initialization. See the AscendCL external API list.
    // Fill in the deviceId based on the actual device.
    int32_t deviceId = 0;
    aclrtStream stream;
    auto ret = Init(deviceId, &stream);
    // Handle check based on the actual situation.
    CHECK_RET(ret == 0, LOG_PRINT("Init acl failed. ERROR: %d\n", ret); return ret);
    // 2. Construct inputs and outputs based on the API definition.
    std::vector<int64_t> tokenXShape = {8, 1, 7168};            // B,S,He
    std::vector<int64_t> weightDqShape = {7168, 1536};          // He,Hcq
    std::vector<int64_t> weightUqQrShape = {1536, 6144};        // Hcq,N*(D+Dr)
    std::vector<int64_t> weightUkShape = {32, 128, 512};        // N,D,Hckv
    std::vector<int64_t> weightDkvKrShape = {7168, 576};        // He,Hckv+Dr
    std::vector<int64_t> rmsnormGammaCqShape = {1536};          // Hcq
    std::vector<int64_t> rmsnormGammaCkvShape = {512};          // Hckv
    std::vector<int64_t> ropeSinShape = {8, 1, 64};             // B,S,Dr
    std::vector<int64_t> ropeCosShape = {8, 1, 64};             // B,S,Dr
    std::vector<int64_t> cacheIndexShape = {8, 1};              // B,S
    std::vector<int64_t> kvCacheShape = {16, 128, 1, 512};      // BlockNum,BlockSize,Nkv,Hckv
    std::vector<int64_t> krCacheShape = {16, 128, 1, 64};       // BlockNum,BlockSize,Nkv,Dr
    std::vector<int64_t> dequantScaleXShape = {8, 1};           // B*S, 1
    std::vector<int64_t> dequantScaleWDqShape = {1, 1536};      // 1, Hcq
    std::vector<int64_t> dequantScaleWUqQrShape = {1, 6144};    // 1, N*(D+Dr)
    std::vector<int64_t> dequantScaleWDkvKrShape = {1,576};     // 1, Hckv+Dr
    std::vector<int64_t> quantScaleCkvShape = {1, 512};         // 1, Hckv
    std::vector<int64_t> smoothScaleCqShape = {1, 1536};        // 1, Hcq
    std::vector<int64_t> queryShape = {8, 1, 32, 512};          // B,S,N,Hckv
    std::vector<int64_t> queryRopeShape = {8, 1, 32, 64};       // B,S,N,Dr
    std::vector<int64_t> dequantScaleQNopeShape = {8, 32, 1};   // B*S, N, 1
    double rmsnormEpsilonCq = 1e-5;
    double rmsnormEpsilonCkv = 1e-5;
    char cacheMode[] = "PA_BSND";

    void* tokenXDeviceAddr = nullptr;
    void* weightDqDeviceAddr = nullptr;
    void* weightUqQrDeviceAddr = nullptr;
    void* weightUkDeviceAddr = nullptr;
    void* weightDkvKrDeviceAddr = nullptr;
    void* rmsnormGammaCqDeviceAddr = nullptr;
    void* rmsnormGammaCkvDeviceAddr = nullptr;
    void* ropeSinDeviceAddr = nullptr;
    void* ropeCosDeviceAddr = nullptr;
    void* cacheIndexDeviceAddr = nullptr;
    void* kvCacheDeviceAddr = nullptr;
    void* krCacheDeviceAddr = nullptr;
    void* dequantScaleXDeviceAddr = nullptr;
    void* dequantScaleWDqDeviceAddr = nullptr;
    void* dequantScaleWUqQrDeviceAddr = nullptr;
    void* dequantScaleWDkvKrDeviceAddr = nullptr;
    void* quantScaleCkvDeviceAddr = nullptr;
    void* smoothScaleCqDeviceAddr = nullptr;
    void* queryDeviceAddr = nullptr;
    void* queryRopeDeviceAddr = nullptr;
    void* dequantScaleQNopeDeviceAddr = nullptr;

    void* tokenXHostAddr = nullptr;
    void* weightDqHostAddr = nullptr;
    void* weightUqQrHostAddr = nullptr;
    void* weightUkHostAddr = nullptr;
    void* weightDkvKrHostAddr = nullptr;
    void* rmsnormGammaCqHostAddr = nullptr;
    void* rmsnormGammaCkvHostAddr = nullptr;
    void* ropeSinHostAddr = nullptr;
    void* ropeCosHostAddr = nullptr;
    void* cacheIndexHostAddr = nullptr;
    void* kvCacheHostAddr = nullptr;
    void* krCacheHostAddr = nullptr;
    void* dequantScaleXHostAddr = nullptr;
    void* dequantScaleWDqHostAddr = nullptr;
    void* dequantScaleWUqQrHostAddr = nullptr;
    void* dequantScaleWDkvKrHostAddr = nullptr;
    void* quantScaleCkvHostAddr = nullptr;
    void* smoothScaleCqHostAddr = nullptr;
    void* queryHostAddr = nullptr;
    void* queryRopeHostAddr = nullptr;
    void* dequantScaleQNopeHostAddr = nullptr;

    aclTensor* tokenX = nullptr;
    aclTensor* weightDq = nullptr;
    aclTensor* weightUqQr = nullptr;
    aclTensor* weightUk = nullptr;
    aclTensor* weightDkvKr = nullptr;
    aclTensor* rmsnormGammaCq = nullptr;
    aclTensor* rmsnormGammaCkv = nullptr;
    aclTensor* ropeSin = nullptr;
    aclTensor* ropeCos = nullptr;
    aclTensor* cacheIndex = nullptr;
    aclTensor* kvCache = nullptr;
    aclTensor* krCache = nullptr;
    aclTensor* dequantScaleX = nullptr;
    aclTensor* dequantScaleWDq = nullptr;
    aclTensor* dequantScaleWUqQr = nullptr;
    aclTensor* dequantScaleWDkvKr = nullptr;
    aclTensor* quantScaleCkv = nullptr;
    aclTensor* smoothScaleCq = nullptr;
    aclTensor* query = nullptr;
    aclTensor* queryRope = nullptr;
    aclTensor* dequantScaleQNope = nullptr;

    // Convert the shapes of three NZ format variables.
    ret = TransToNZShape(weightDqShape, sizeof(int8_t));
    CHECK_RET(ret == 0, LOG_PRINT("trans NZ shape failed. \n"); return ret);
    ret = TransToNZShape(weightUqQrShape, sizeof(int8_t));
    CHECK_RET(ret == 0, LOG_PRINT("trans NZ shape failed. \n"); return ret);
    ret = TransToNZShape(weightDkvKrShape, sizeof(int8_t));
    CHECK_RET(ret == 0, LOG_PRINT("trans NZ shape failed. \n"); return ret);

    // Create the tokenX aclTensor.
    ret = CreateAclTensorND(tokenXShape, &tokenXDeviceAddr, &tokenXHostAddr, aclDataType::ACL_INT8, &tokenX);
    CHECK_RET(ret == ACL_SUCCESS, return ret);
    // Create the weightDq aclTensor.
    ret = CreateAclTensorNZ(weightDqShape, &weightDqDeviceAddr, &weightDqHostAddr, aclDataType::ACL_INT8, &weightDq);
    CHECK_RET(ret == ACL_SUCCESS, return ret);
    // Create the weightUqQr aclTensor.
    ret = CreateAclTensorNZ(weightUqQrShape, &weightUqQrDeviceAddr, &weightUqQrHostAddr, aclDataType::ACL_INT8, &weightUqQr);
    CHECK_RET(ret == ACL_SUCCESS, return ret);
    // Create the weightUk aclTensor.
    ret = CreateAclTensorND(weightUkShape, &weightUkDeviceAddr, &weightUkHostAddr, aclDataType::ACL_BF16, &weightUk);
    CHECK_RET(ret == ACL_SUCCESS, return ret);
    // Create the weightDkvKr aclTensor.
    ret = CreateAclTensorNZ(weightDkvKrShape, &weightDkvKrDeviceAddr, &weightDkvKrHostAddr, aclDataType::ACL_INT8, &weightDkvKr);
    CHECK_RET(ret == ACL_SUCCESS, return ret);
    // Create the rmsnormGammaCq aclTensor.
    ret = CreateAclTensorND(rmsnormGammaCqShape, &rmsnormGammaCqDeviceAddr, &rmsnormGammaCqHostAddr, aclDataType::ACL_BF16, &rmsnormGammaCq);
    CHECK_RET(ret == ACL_SUCCESS, return ret);
    // Create the rmsnormGammaCkv aclTensor.
    ret = CreateAclTensorND(rmsnormGammaCkvShape, &rmsnormGammaCkvDeviceAddr, &rmsnormGammaCkvHostAddr, aclDataType::ACL_BF16, &rmsnormGammaCkv);
    CHECK_RET(ret == ACL_SUCCESS, return ret);
    // Create the ropeSin aclTensor.
    ret = CreateAclTensorND(ropeSinShape, &ropeSinDeviceAddr, &ropeSinHostAddr, aclDataType::ACL_BF16, &ropeSin);
    CHECK_RET(ret == ACL_SUCCESS, return ret);
    // Create the ropeCos aclTensor.
    ret = CreateAclTensorND(ropeCosShape, &ropeCosDeviceAddr, &ropeCosHostAddr, aclDataType::ACL_BF16, &ropeCos);
    CHECK_RET(ret == ACL_SUCCESS, return ret);
    // Create the cacheIndex aclTensor.
    ret = CreateAclTensorND(cacheIndexShape, &cacheIndexDeviceAddr, &cacheIndexHostAddr, aclDataType::ACL_INT64, &cacheIndex);
    CHECK_RET(ret == ACL_SUCCESS, return ret);
    // Create the kvCache aclTensor.
    ret = CreateAclTensorND(kvCacheShape, &kvCacheDeviceAddr, &kvCacheHostAddr, aclDataType::ACL_INT8, &kvCache);
    CHECK_RET(ret == ACL_SUCCESS, return ret);
    // Create the krCache aclTensor.
    ret = CreateAclTensorND(krCacheShape, &krCacheDeviceAddr, &krCacheHostAddr, aclDataType::ACL_BF16, &krCache);
    CHECK_RET(ret == ACL_SUCCESS, return ret);
    // Create the dequantScaleX aclTensor.
    ret = CreateAclTensorND(dequantScaleXShape, &dequantScaleXDeviceAddr, &dequantScaleXHostAddr, aclDataType::ACL_FLOAT, &dequantScaleX);
    CHECK_RET(ret == ACL_SUCCESS, return ret);
    // Create the dequantScaleWDq aclTensor.
    ret = CreateAclTensorND(dequantScaleWDqShape, &dequantScaleWDqDeviceAddr, &dequantScaleWDqHostAddr, aclDataType::ACL_FLOAT, &dequantScaleWDq);
    CHECK_RET(ret == ACL_SUCCESS, return ret);
    // Create the dequantScaleWUqQr aclTensor.
    ret = CreateAclTensorND(dequantScaleWUqQrShape, &dequantScaleWUqQrDeviceAddr, &dequantScaleWUqQrHostAddr, aclDataType::ACL_FLOAT, &dequantScaleWUqQr);
    CHECK_RET(ret == ACL_SUCCESS, return ret);
    // Create the dequantScaleWDkvKr aclTensor.
    ret = CreateAclTensorND(dequantScaleWDkvKrShape, &dequantScaleWDkvKrDeviceAddr, &dequantScaleWDkvKrHostAddr, aclDataType::ACL_FLOAT, &dequantScaleWDkvKr);
    CHECK_RET(ret == ACL_SUCCESS, return ret);
    // Create the quantScaleCkv aclTensor.
    ret = CreateAclTensorND(quantScaleCkvShape, &quantScaleCkvDeviceAddr, &quantScaleCkvHostAddr, aclDataType::ACL_FLOAT, &quantScaleCkv);
    CHECK_RET(ret == ACL_SUCCESS, return ret);
    // Create the smoothScaleCq aclTensor.
    ret = CreateAclTensorND(smoothScaleCqShape, &smoothScaleCqDeviceAddr, &smoothScaleCqHostAddr, aclDataType::ACL_FLOAT, &smoothScaleCq);
    CHECK_RET(ret == ACL_SUCCESS, return ret);
    // Create the query aclTensor.
    ret = CreateAclTensorND(queryShape, &queryDeviceAddr, &queryHostAddr, aclDataType::ACL_INT8, &query);
    CHECK_RET(ret == ACL_SUCCESS, return ret);
    // Create the queryRope aclTensor.
    ret = CreateAclTensorND(queryRopeShape, &queryRopeDeviceAddr, &queryRopeHostAddr, aclDataType::ACL_BF16, &queryRope);
    CHECK_RET(ret == ACL_SUCCESS, return ret);
    // Create the dequantScaleQNope aclTensor.
    ret = CreateAclTensorND(dequantScaleQNopeShape, &dequantScaleQNopeDeviceAddr, &dequantScaleQNopeHostAddr, aclDataType::ACL_FLOAT, &dequantScaleQNope);
    CHECK_RET(ret == ACL_SUCCESS, return ret);

    // 3. Call the CANN operator library API. Replace it with the specific API.
    uint64_t workspaceSize = 0;
    aclOpExecutor* executor = nullptr;
    // Call the first-phase API of aclnnMlaPrologV2WeightNz.
    ret = aclnnMlaPrologV2WeightNzGetWorkspaceSize(tokenX, weightDq, weightUqQr, weightUk, weightDkvKr, rmsnormGammaCq, rmsnormGammaCkv, ropeSin, ropeCos, cacheIndex, kvCache, krCache,
      dequantScaleX, dequantScaleWDq, dequantScaleWUqQr, dequantScaleWDkvKr, quantScaleCkv, nullptr, smoothScaleCq, rmsnormEpsilonCq, rmsnormEpsilonCkv, cacheMode,
      query, queryRope, dequantScaleQNope, &workspaceSize, &executor);
    CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("aclnnMlaPrologV2WeightNzGetWorkspaceSize failed. ERROR: %d\n", ret); return ret);
    // Allocate device memory based on the workspaceSize computed by the first-phase API.
    void* workspaceAddr = nullptr;
    if (workspaceSize > 0) {
        ret = aclrtMalloc(&workspaceAddr, workspaceSize, ACL_MEM_MALLOC_HUGE_FIRST);
        CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("allocate workspace failed. ERROR: %d\n", ret); return ret;);
    }
    // Call the second-phase API of aclnnMlaPrologV2WeightNz.
    ret = aclnnMlaPrologV2WeightNz(workspaceAddr, workspaceSize, executor, stream);
    CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("aclnnMlaPrologV2WeightNz failed. ERROR: %d\n", ret); return ret);

    // 4. Boilerplate: synchronize and wait for task execution to complete.
    ret = aclrtSynchronizeStream(stream);
    CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("aclrtSynchronizeStream failed. ERROR: %d\n", ret); return ret);

    // 5. Obtain the output value and copy the result from device memory to host memory. Modify based on the specific API definition.
    auto size = GetShapeSize(queryShape);
    std::vector<float> resultData(size, 0);
    ret = aclrtMemcpy(resultData.data(), resultData.size() * sizeof(resultData[0]), queryDeviceAddr, size * sizeof(float),
                      ACL_MEMCPY_DEVICE_TO_HOST);
    CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("copy result from device to host failed. ERROR: %d\n", ret); return ret);

    // 6. Release aclTensor and aclScalar. Modify based on the specific API definition.
    aclDestroyTensor(tokenX);
    aclDestroyTensor(weightDq);
    aclDestroyTensor(weightUqQr);
    aclDestroyTensor(weightUk);
    aclDestroyTensor(weightDkvKr);
    aclDestroyTensor(rmsnormGammaCq);
    aclDestroyTensor(rmsnormGammaCkv);
    aclDestroyTensor(ropeSin);
    aclDestroyTensor(ropeCos);
    aclDestroyTensor(cacheIndex);
    aclDestroyTensor(kvCache);
    aclDestroyTensor(krCache);
    aclDestroyTensor(dequantScaleX);
    aclDestroyTensor(dequantScaleWDq);
    aclDestroyTensor(dequantScaleWUqQr);
    aclDestroyTensor(dequantScaleWDkvKr);
    aclDestroyTensor(quantScaleCkv);
    aclDestroyTensor(smoothScaleCq);
    aclDestroyTensor(query);
    aclDestroyTensor(queryRope);
    aclDestroyTensor(dequantScaleQNope);

    // 7. Release device resources.
    aclrtFree(tokenXDeviceAddr);
    aclrtFree(weightDqDeviceAddr);
    aclrtFree(weightUqQrDeviceAddr);
    aclrtFree(weightUkDeviceAddr);
    aclrtFree(weightDkvKrDeviceAddr);
    aclrtFree(rmsnormGammaCqDeviceAddr);
    aclrtFree(rmsnormGammaCkvDeviceAddr);
    aclrtFree(ropeSinDeviceAddr);
    aclrtFree(ropeCosDeviceAddr);
    aclrtFree(cacheIndexDeviceAddr);
    aclrtFree(kvCacheDeviceAddr);
    aclrtFree(krCacheDeviceAddr);
    aclrtFree(dequantScaleXDeviceAddr);
    aclrtFree(dequantScaleWDqDeviceAddr);
    aclrtFree(dequantScaleWUqQrDeviceAddr);
    aclrtFree(dequantScaleWDkvKrDeviceAddr);
    aclrtFree(quantScaleCkvDeviceAddr);
    aclrtFree(smoothScaleCqDeviceAddr);
    aclrtFree(queryDeviceAddr);
    aclrtFree(queryRopeDeviceAddr);
    aclrtFree(dequantScaleQNopeDeviceAddr);

    // 8. Release host resources.
    aclrtFree(tokenXHostAddr);
    aclrtFree(weightDqHostAddr);
    aclrtFree(weightUqQrHostAddr);
    aclrtFree(weightUkHostAddr);
    aclrtFree(weightDkvKrHostAddr);
    aclrtFree(rmsnormGammaCqHostAddr);
    aclrtFree(rmsnormGammaCkvHostAddr);
    aclrtFree(ropeSinHostAddr);
    aclrtFree(ropeCosHostAddr);
    aclrtFree(cacheIndexHostAddr);
    aclrtFree(kvCacheHostAddr);
    aclrtFree(krCacheHostAddr);
    aclrtFree(dequantScaleXHostAddr);
    aclrtFree(dequantScaleWDqHostAddr);
    aclrtFree(dequantScaleWUqQrHostAddr);
    aclrtFree(dequantScaleWDkvKrHostAddr);
    aclrtFree(quantScaleCkvHostAddr);
    aclrtFree(smoothScaleCqHostAddr);
    aclrtFree(queryHostAddr);
    aclrtFree(queryRopeHostAddr);
    aclrtFree(dequantScaleQNopeHostAddr);

    if (workspaceSize > 0) {
      aclrtFree(workspaceAddr);
    }
    aclrtDestroyStream(stream);
    aclrtResetDevice(deviceId);
    aclFinalize();

    return 0;
}

The following is a sample code for A5, for reference only.

#include <iostream>
#include <vector>
#include <cstdint>
#include "acl/acl.h"
#include "aclnnop/aclnn_mla_prolog_v2_weight_nz.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)

int64_t GetShapeSize(const std::vector<int64_t>& shape) {
    int64_t shape_size = 1;
    for (auto i : shape) {
        shape_size *= i;
    }
    return shape_size;
}

int Init(int32_t deviceId, aclrtStream* stream) {
    // Boilerplate for resource initialization.
    auto ret = aclInit(nullptr);
    CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("aclInit failed. ERROR: %d\n", ret); return ret);
    ret = aclrtSetDevice(deviceId);
    CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("aclrtSetDevice failed. ERROR: %d\n", ret); return ret);
    ret = aclrtCreateStream(stream);
    CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("aclrtCreateStream failed. ERROR: %d\n", ret); return ret);
    return 0;
}

template <typename T>
int CreateAclTensorND(const std::vector<T>& shape, void** deviceAddr, void** hostAddr,
                    aclDataType dataType, aclTensor** tensor) {
    auto size = GetShapeSize(shape) * sizeof(T);
    // Call aclrtMalloc to apply for device memory.
    auto ret = aclrtMalloc(deviceAddr, size, ACL_MEM_MALLOC_HUGE_FIRST);
    CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("aclrtMalloc failed. ERROR: %d\n", ret); return ret);
    // Call aclrtMalloc to apply for host memory.
    ret = aclrtMalloc(hostAddr, size, ACL_MEM_MALLOC_HUGE_FIRST);
    CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("aclrtMalloc failed. ERROR: %d\n", ret); return ret);
    // Call the aclCreateTensor API to create an aclTensor.
    *tensor = aclCreateTensor(shape.data(), shape.size(), dataType, nullptr, 0, aclFormat::ACL_FORMAT_ND,
                              shape.data(), shape.size(), *deviceAddr);
    // Call aclrtMemcpy to copy data from the host to the device memory.
    ret = aclrtMemcpy(*deviceAddr, size, *hostAddr, GetShapeSize(shape)*aclDataTypeSize(dataType), ACL_MEMCPY_HOST_TO_DEVICE);
    CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("aclrtMemcpy failed. ERROR: %d\n", ret); return ret);
    return 0;
}

template <typename T>
int CreateAclTensorNZ(const std::vector<T>& shape, void** deviceAddr, void** hostAddr,
                    aclDataType dataType, aclTensor** tensor) {
    auto size = GetShapeSize(shape) * sizeof(T);
    // Call aclrtMalloc to apply for device memory.
    auto ret = aclrtMalloc(deviceAddr, size, ACL_MEM_MALLOC_HUGE_FIRST);
    CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("aclrtMalloc failed. ERROR: %d\n", ret); return ret);
    // Call aclrtMalloc to allocate host-side memory.
    ret = aclrtMalloc(hostAddr, size, ACL_MEM_MALLOC_HUGE_FIRST);
    CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("aclrtMalloc failed. ERROR: %d\n", ret); return ret);
    // Call the aclCreateTensor API to create an aclTensor.
    *tensor = aclCreateTensor(shape.data(), shape.size(), dataType, nullptr, 0, aclFormat::ACL_FORMAT_FRACTAL_NZ,
                              shape.data(), shape.size(), *deviceAddr);
    // Call aclrtMemcpy to copy host-side data to device-side memory.
    ret = aclrtMemcpy(*deviceAddr, size, *hostAddr, GetShapeSize(shape)*aclDataTypeSize(dataType), ACL_MEMCPY_HOST_TO_DEVICE);
    CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("aclrtMemcpy failed. ERROR: %d\n", ret); return ret);
    return 0;
}

int TransToNZShape(std::vector<int64_t> &shapeND, size_t typeSize) {
    if (typeSize == static_cast<size_t>(0)) {
      return 0;
    }
    int64_t h = shapeND[0];
    int64_t w = shapeND[1];
    int64_t h0 = static_cast<int64_t>(16);
    int64_t w0 = static_cast<int64_t>(32) / static_cast<int64_t>(typeSize);
    int64_t h1 = h / h0;
    int64_t w1 = w / w0;
    shapeND[0] = w1;
    shapeND[1] = h1;
    shapeND.emplace_back(h0);
    shapeND.emplace_back(w0);
    return 0;
}

int main() {
    // 1. Boilerplate: device/stream initialization. See the AscendCL external API list.
    // Fill in deviceId based on your actual device.
    int32_t deviceId = 0;
    aclrtStream stream;
    auto ret = Init(deviceId, &stream);
    // Handle the check based on your own needs.
    CHECK_RET(ret == 0, LOG_PRINT("Init acl failed. ERROR: %d\n", ret); return ret);
    // 2. Construct the input and output based on the API definition.
    std::vector<int64_t> tokenXShape = {8, 1, 7168};            // B,S,He
    std::vector<int64_t> weightDqShape = {7168, 1536};          // He,Hcq
    std::vector<int64_t> weightUqQrShape = {1536, 6144};        // Hcq,N*(D+Dr)
    std::vector<int64_t> weightUkShape = {32, 128, 512};        // N,D,Hckv
    std::vector<int64_t> weightDkvKrShape = {7168, 576};        // He,Hckv+Dr
    std::vector<int64_t> rmsnormGammaCqShape = {1536};          // Hcq
    std::vector<int64_t> rmsnormGammaCkvShape = {512};          // Hckv
    std::vector<int64_t> ropeSinShape = {8, 1, 64};             // B,S,Dr
    std::vector<int64_t> ropeCosShape = {8, 1, 64};             // B,S,Dr
    std::vector<int64_t> cacheIndexShape = {8, 1};              // B,S
    std::vector<int64_t> kvCacheShape = {16, 128, 1, 512};      // BolckNum,BlockSize,Nkv,Hckv
    std::vector<int64_t> krCacheShape = {16, 128, 1, 64};       // BolckNum,BlockSize,Nkv,Dr
    std::vector<int64_t> queryShape = {8, 1, 32, 512};          // B,S,N,Hckv
    std::vector<int64_t> queryRopeShape = {8, 1, 32, 64};       // B,S,N,Dr
    double rmsnormEpsilonCq = 1e-5;
    double rmsnormEpsilonCkv = 1e-5;
    char cacheMode[] = "PA_BSND";

    void* tokenXDeviceAddr = nullptr;
    void* weightDqDeviceAddr = nullptr;
    void* weightUqQrDeviceAddr = nullptr;
    void* weightUkDeviceAddr = nullptr;
    void* weightDkvKrDeviceAddr = nullptr;
    void* rmsnormGammaCqDeviceAddr = nullptr;
    void* rmsnormGammaCkvDeviceAddr = nullptr;
    void* ropeSinDeviceAddr = nullptr;
    void* ropeCosDeviceAddr = nullptr;
    void* cacheIndexDeviceAddr = nullptr;
    void* kvCacheDeviceAddr = nullptr;
    void* krCacheDeviceAddr = nullptr;
    void* queryDeviceAddr = nullptr;
    void* queryRopeDeviceAddr = nullptr;

    void* tokenXHostAddr = nullptr;
    void* weightDqHostAddr = nullptr;
    void* weightUqQrHostAddr = nullptr;
    void* weightUkHostAddr = nullptr;
    void* weightDkvKrHostAddr = nullptr;
    void* rmsnormGammaCqHostAddr = nullptr;
    void* rmsnormGammaCkvHostAddr = nullptr;
    void* ropeSinHostAddr = nullptr;
    void* ropeCosHostAddr = nullptr;
    void* cacheIndexHostAddr = nullptr;
    void* kvCacheHostAddr = nullptr;
    void* krCacheHostAddr = nullptr;
    void* queryHostAddr = nullptr;
    void* queryRopeHostAddr = nullptr;

    aclTensor* tokenX = nullptr;
    aclTensor* weightDq = nullptr;
    aclTensor* weightUqQr = nullptr;
    aclTensor* weightUk = nullptr;
    aclTensor* weightDkvKr = nullptr;
    aclTensor* rmsnormGammaCq = nullptr;
    aclTensor* rmsnormGammaCkv = nullptr;
    aclTensor* ropeSin = nullptr;
    aclTensor* ropeCos = nullptr;
    aclTensor* cacheIndex = nullptr;
    aclTensor* kvCache = nullptr;
    aclTensor* krCache = nullptr;
    aclTensor* query = nullptr;
    aclTensor* queryRope = nullptr;

    // Convert the shapes of the three NZ-format variables.
    constexpr size_t EXAMPLE_INT8_SIZE = sizeof(int8_t);
    constexpr size_t EXAMPLE_BFLOAT16_SIZE = sizeof(int16_t);
    ret = TransToNZShape(weightDqShape, EXAMPLE_BFLOAT16_SIZE);
    CHECK_RET(ret == 0, LOG_PRINT("trans NZ shape failed.\n"); return ret);
    ret = TransToNZShape(weightUqQrShape, EXAMPLE_BFLOAT16_SIZE);
    CHECK_RET(ret == 0, LOG_PRINT("trans NZ shape failed.\n"); return ret);
    ret = TransToNZShape(weightDkvKrShape, EXAMPLE_BFLOAT16_SIZE);
    CHECK_RET(ret == 0, LOG_PRINT("trans NZ shape failed.\n"); return ret);

    // Create the tokenX aclTensor.
    ret = CreateAclTensorND(tokenXShape, &tokenXDeviceAddr, &tokenXHostAddr, aclDataType::ACL_BF16, &tokenX);
    CHECK_RET(ret == ACL_SUCCESS, return ret);
    // Create the weightDq aclTensor.
    ret = CreateAclTensorNZ(weightDqShape, &weightDqDeviceAddr, &weightDqHostAddr, aclDataType::ACL_BF16, &weightDq);
    CHECK_RET(ret == ACL_SUCCESS, return ret);
    // Create weightUqQr aclTensor.
    ret = CreateAclTensorNZ(weightUqQrShape, &weightUqQrDeviceAddr, &weightUqQrHostAddr, aclDataType::ACL_BF16, &weightUqQr);
    CHECK_RET(ret == ACL_SUCCESS, return ret);
    // Create weightUk aclTensor.
    ret = CreateAclTensorND(weightUkShape, &weightUkDeviceAddr, &weightUkHostAddr, aclDataType::ACL_BF16, &weightUk);
    CHECK_RET(ret == ACL_SUCCESS, return ret);
    // Create weightDkvKr aclTensor.
    ret = CreateAclTensorNZ(weightDkvKrShape, &weightDkvKrDeviceAddr, &weightDkvKrHostAddr, aclDataType::ACL_BF16, &weightDkvKr);
    CHECK_RET(ret == ACL_SUCCESS, return ret);
    // Create ropeSin aclTensor.
    ret = CreateAclTensorND(ropeSinShape, &ropeSinDeviceAddr, &ropeSinHostAddr, aclDataType::ACL_BF16, &ropeSin);
    CHECK_RET(ret == ACL_SUCCESS, return ret);
    // Create ropeCos aclTensor.
    ret = CreateAclTensorND(ropeCosShape, &ropeCosDeviceAddr, &ropeCosHostAddr, aclDataType::ACL_BF16, &ropeCos);
    CHECK_RET(ret == ACL_SUCCESS, return ret);
    // Create the rmsnormGammaCq aclTensor.
    ret = CreateAclTensorND(rmsnormGammaCqShape, &rmsnormGammaCqDeviceAddr, &rmsnormGammaCqHostAddr, aclDataType::ACL_BF16, &rmsnormGammaCq);
    CHECK_RET(ret == ACL_SUCCESS, return ret);
    // Create the rmsnormGammaCkv aclTensor.
    ret = CreateAclTensorND(rmsnormGammaCkvShape, &rmsnormGammaCkvDeviceAddr, &rmsnormGammaCkvHostAddr, aclDataType::ACL_BF16, &rmsnormGammaCkv);
    CHECK_RET(ret == ACL_SUCCESS, return ret);
    // Create the cacheIndex aclTensor.
    ret = CreateAclTensorND(cacheIndexShape, &cacheIndexDeviceAddr, &cacheIndexHostAddr, aclDataType::ACL_INT64, &cacheIndex);
    CHECK_RET(ret == ACL_SUCCESS, return ret);
    // Create the kvCache aclTensor.
    ret = CreateAclTensorND(kvCacheShape, &kvCacheDeviceAddr, &kvCacheHostAddr, aclDataType::ACL_BF16, &kvCache);
    CHECK_RET(ret == ACL_SUCCESS, return ret);
    // Create the krCache aclTensor.
    ret = CreateAclTensorND(krCacheShape, &krCacheDeviceAddr, &krCacheHostAddr, aclDataType::ACL_BF16, &krCache);
    CHECK_RET(ret == ACL_SUCCESS, return ret);
    // Create the query aclTensor.
    ret = CreateAclTensorND(queryShape, &queryDeviceAddr, &queryHostAddr, aclDataType::ACL_BF16, &query);
    CHECK_RET(ret == ACL_SUCCESS, return ret);
    // Create the queryRope aclTensor.
    ret = CreateAclTensorND(queryRopeShape, &queryRopeDeviceAddr, &queryRopeHostAddr, aclDataType::ACL_BF16, &queryRope);
    CHECK_RET(ret == ACL_SUCCESS, return ret);

    // 3. Call the CANN operator library API. Replace with the specific API.
    uint64_t workspaceSize = 0;
    aclOpExecutor* executor = nullptr;
    // Call the first-phase API of aclnnMlaPrologV2WeightNz.
    ret = aclnnMlaPrologV2WeightNzGetWorkspaceSize(tokenX, weightDq, weightUqQr, weightUk, weightDkvKr, rmsnormGammaCq, rmsnormGammaCkv, ropeSin, ropeCos, cacheIndex, kvCache, krCache, nullptr, nullptr, nullptr, nullptr, nullptr, nullptr, nullptr, rmsnormEpsilonCq, rmsnormEpsilonCkv, cacheMode, query, queryRope, nullptr, &workspaceSize, &executor);
    CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("aclnnMlaPrologV2WeightNzGetWorkspaceSize failed. ERROR: %d\n", ret); return ret);
    // Apply for device memory based on the workspaceSize computed by the first-phase API.
    void* workspaceAddr = nullptr;
    if (workspaceSize > static_cast<uint64_t>(0)) {
        ret = aclrtMalloc(&workspaceAddr, workspaceSize, ACL_MEM_MALLOC_HUGE_FIRST);
        CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("allocate workspace failed. ERROR: %d\n", ret); return ret;);
    }
    // Call the second-phase API of aclnnMlaPrologV2WeightNz.
    ret = aclnnMlaPrologV2WeightNz(workspaceAddr, workspaceSize, executor, stream);
    CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("aclnnMlaPrologV2WeightNz failed. ERROR: %d\n", ret); return ret);

    // 4. Boilerplate: synchronize and wait for task execution to complete.
    ret = aclrtSynchronizeStream(stream);
    CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("aclrtSynchronizeStream failed. ERROR: %d\n", ret); return ret);

    // 5. Obtain the output value and copy the result from the device memory to the host memory. Modify based on the specific API definition.
    auto size = GetShapeSize(queryShape);
    std::vector<float> resultData(size, 0);
    ret = aclrtMemcpy(resultData.data(), resultData.size() * sizeof(resultData[0]), queryDeviceAddr, size * sizeof(float),
                      ACL_MEMCPY_DEVICE_TO_HOST);
    CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("copy result from device to host failed. ERROR: %d\n", ret); return ret);
    for (int64_t i = 0; i < size; i++) {
      LOG_PRINT("result[%ld] is: %f\n", i, resultData[i]);
    }
    // 6. Release aclTensor and aclScalar. Modify based on the specific API definition.
    aclDestroyTensor(tokenX);
    aclDestroyTensor(weightDq);
    aclDestroyTensor(weightUqQr);
    aclDestroyTensor(weightUk);
    aclDestroyTensor(weightDkvKr);
    aclDestroyTensor(rmsnormGammaCq);
    aclDestroyTensor(rmsnormGammaCkv);
    aclDestroyTensor(ropeSin);
    aclDestroyTensor(ropeCos);
    aclDestroyTensor(cacheIndex);
    aclDestroyTensor(kvCache);
    aclDestroyTensor(krCache);
    aclDestroyTensor(query);
    aclDestroyTensor(queryRope);

    // 7. Release device resources.
    aclrtFree(tokenXDeviceAddr);
    aclrtFree(weightDqDeviceAddr);
    aclrtFree(weightUqQrDeviceAddr);
    aclrtFree(weightUkDeviceAddr);
    aclrtFree(weightDkvKrDeviceAddr);
    aclrtFree(rmsnormGammaCqDeviceAddr);
    aclrtFree(rmsnormGammaCkvDeviceAddr);
    aclrtFree(ropeSinDeviceAddr);
    aclrtFree(ropeCosDeviceAddr);
    aclrtFree(cacheIndexDeviceAddr);
    aclrtFree(kvCacheDeviceAddr);
    aclrtFree(krCacheDeviceAddr);
    aclrtFree(queryDeviceAddr);
    aclrtFree(queryRopeDeviceAddr);

    // 8. Release host resources.
    aclrtFree(tokenXHostAddr);
    aclrtFree(weightDqHostAddr);
    aclrtFree(weightUqQrHostAddr);
    aclrtFree(weightUkHostAddr);
    aclrtFree(weightDkvKrHostAddr);
    aclrtFree(rmsnormGammaCqHostAddr);
    aclrtFree(rmsnormGammaCkvHostAddr);
    aclrtFree(ropeSinHostAddr);
    aclrtFree(ropeCosHostAddr);
    aclrtFree(cacheIndexHostAddr);
    aclrtFree(kvCacheHostAddr);
    aclrtFree(krCacheHostAddr);
    aclrtFree(queryHostAddr);
    aclrtFree(queryRopeHostAddr);

    if (workspaceSize > static_cast<uint64_t>(0)) {
      aclrtFree(workspaceAddr);
    }
    aclrtDestroyStream(stream);
    aclrtResetDevice(deviceId);
    aclFinalize();

    return 0;
}