aclnnMlaPrologV3WeightNz

📄 View Source Code

Applicable Products

Function

• Feature Update: (Differences from aclnnMlaPrologV2weightNz)

  • Added scale correction factors for Query and Key, corresponding to qcQrScale ($\alpha_q$) and kcScale ($\alpha_{kv}$), respectively.

  • Added optional input parameters (such as actualSeqLenOptional, kNopeClipAlphaOptional, queryNormFlag, weightQuantMode, kvCacheQuantMode, queryQuantMode, ckvkrRepoMode, quantScaleRepoMode, tileSize, queryNormOutOptional, and dequantScaleQNormOptional), and changed cache_mode from mandatory to optional.

  • Adjusted the name and position of the cacheIndex parameter, corresponding to the current cacheIndexOptional.

• Description: In inference scenarios, this API performs Multi-Head Latent Attention pre-processing. 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}$, undergoes two upsampling operations, and is then multiplied by the Query scale correction factor $\alpha_q$ 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, and then multiplied by the Key scale correction factor $\alpha_{kv}$ and passed into a Cache to obtain $k^C$.

  • The fourth path: the input $x$ is multiplied by $W^{KR}$, then undergoes rotary positional encoding and is passed into another cache to obtain $k^R$.

  • The fifth path: the quantization parameters obtained after $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 calculation formula of Query, including downsampling, RmsNorm, and two upsampling operations

  $$c^Q = \alpha_q\cdot\mathrm{RmsNorm}(x \cdot W^{DQ})$$ $$q^C = c^Q \cdot W^{UQ}$$ $$q^N = q^C \cdot W^{UK}$$

  Applies ROPE rotary position encoding to Query.

  $$q^R = \mathrm{ROPE}(c^Q \cdot W^{QR})$$

  Key calculation formula, including downsampling and RmsNorm, and stores the result in cache.

  $$c^{KV} = \alpha_{kv}\cdot\mathrm{RmsNorm}(x \cdot W^{DKV})$$ $$k^C = \mathrm{Cache}(c^{KV})$$

  Applies ROPE rotary position encoding to the Key and stores the result in the cache.

  $$k^R = \mathrm{Cache}(\mathrm{ROPE}(x \cdot W^{KR}))$$

  Dequant Scale Query Nope formula

  $$\mathrm{dequantScaleQNope} = {\mathrm{RowMax}(\mathrm{abs}(q^{N})) / 127}$$ $$q^{N} = {\mathrm{round}(q^{N} / \mathrm{dequantScaleQNope})}$$

Function Prototype

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

aclnnStatus aclnnMlaPrologV3WeightNzGetWorkspaceSize(
    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,
    aclTensor       *kvCacheRef,
    aclTensor       *krCacheRef,
    const aclTensor *cacheIndexOptional,
    const aclTensor *dequantScaleXOptional,
    const aclTensor *dequantScaleWDqOptional,
    const aclTensor *dequantScaleWUqQrOptional,
    const aclTensor *dequantScaleWDkvKrOptional,
    const aclTensor *quantScaleCkvOptional,
    const aclTensor *quantScaleCkrOptional,
    const aclTensor *smoothScalesCqOptional,
    const aclTensor *actualSeqLenOptional,
    const aclTensor *kNopeClipAlphaOptional,
    double           rmsnormEpsilonCq,
    double           rmsnormEpsilonCkv,
    char            *cacheModeOptional,
    int64_t          weightQuantMode,
    int64_t          kvCacheQuantMode,
    int64_t          queryQuantMode,
    int64_t          ckvkrRepoMode,
    int64_t          quantScaleRepoMode,
    int64_t          tileSize,
    double           qcQrScale,
    double           kcScale,
    const aclTensor *queryOut,
    const aclTensor *queryRopeOut,
    const aclTensor *dequantScaleQNopeOutOptional,
    const aclTensor *queryNormOutOptional,
    const aclTensor *dequantScaleQNormOutOptional,
    uint64_t        *workspaceSize,
    aclOpExecutor  **executor)

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

aclnnMlaPrologV3WeightNzGetWorkspaceSize

• Parameters

  Parameter	Input or Output	Description	Instruction	Data Type	Supported Data Types for Ascend 950PR/Ascend 950DT	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	BFLOAT16, FLOAT8_E4M3FN, INT8, HIFLOAT8	ND	- BS concatenated axis: (T, He) - BS non-concatenated axis: (B, S, He)	×
  weightDq	Input	Downsampling weight matrix $W^{DQ}$ used to compute Query in the formula. Dimension representation without transposition: (k, n).	- Does not support empty tensor.	BFLOAT16, INT8	BFLOAT16, FLOAT8_E4M3FN, INT8, HIFLOAT8	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. Dimension representation without transposition: (k, n).	- Does not support empty tensor.	BFLOAT16, INT8	BFLOAT16, INT8, FLOAT8_E4M3FN, HIFLOAT8	FRACTAL_NZ	(Hcq, N*(D+Dr))	×
  weightUk	Input	Upsampling weight $W^{UK}$ used to compute Key in the formula.	- Does not support empty tensor.	BFLOAT16	BFLOAT16	ND	(N, D, Hckv)	×
  weightDkvKr	Input	Downsampling weight matrix $W^{DKV}$ and positional encoding weight matrix $W^{KR}$ used to compute Key in the formula. Dimension representation without transposition: (k, n).	- Does not support empty tensor.	BFLOAT16, INT8	BFLOAT16, FLOAT8_E4M3FN, INT8, HIFLOAT8	FRACTAL_NZ	(He, Hckv+Dr)	×
  rmsnormGammaCq	Input	The $\gamma$ parameter in the RmsNorm formula for computing $c^Q$.	- Does not support empty tensor.	BFLOAT16	BFLOAT16	ND	(Hcq)	×
  rmsnormGammaCkv	Input	The $\gamma$ parameter in the RmsNorm formula for computing $c^{KV}$.	- Does not support empty tensor.	BFLOAT16	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	BFLOAT16	ND	- BS concatenated axis: (T, Dr) - BS 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	BFLOAT16	ND	- BS concatenated axis: (T, Dr) - BS non-concatenated axis: (B, S, Dr)	×
  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	BFLOAT16, INT8, FLOAT8_E4M3FN, HIFLOAT8	ND	- CacheMode="PA_BSND"/"PA_NZ"/"PA_BLK_BSND"/"PA_BLK_NZ": (BlockNum, BlockSize, Nkv, Dtile) - CacheMode="BSND": (B, S, Nkv, Dtile) - CacheMode="TND": (T, Nkv, Dtile)	×
  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	BFLOAT16, INT8	ND	- CacheMode="PA_BSND"/"PA_NZ"/"PA_BLK_BSND"/"PA_BLK_NZ": (BlockNum, BlockSize, Nkv, Dr) - CacheMode="BSND": (B, S, Nkv, Dr) - CacheMode="TND": (T, Nkv, Dr) - When ckvkrRepoMode=1: The dimension must contain 0, and (0) is supported.	×
  cacheIndexOptional	Input	Stores the indices of kvCache and krCache.	- Supports empty tensor with B=0, S=0, T=0. - cacheMode="PA_BSND"/"PA_NZ": The value range must be within [0, BlockNum*BlockSize). - cacheMode="PA_BLK_BSND"/"PA_BLK_NZ": The value range must be within [0, BlockNum). - cacheMode="TND"/"BSND": nullptr.	INT64	INT64	ND	- CacheMode="PA_BSND"/"PA_NZ": 1. BS concatenated axis: (T) 2. BS non-concatenated axis: (B, S) - CacheMode="PA_BLK_BSND"/"PA_BLK_NZ": 1. BS concatenated axis: (Sum(Ceil(S_i/BlockSize))), where S_i is the length of S in each batch. 2. BS non-concatenated axis: (B, Ceil(S/BlockSize)) - CacheMode="TND"/"BSND": nullptr.	×
  dequantScaleXOptional	Input	Dequantization parameter of tokenX.	- Supports empty tensor with B=0, S=0, T=0 (must be passed in weightQuantMode=2/3/4/5 scenarios).	FLOAT	FLOAT8_E8M0, FLOAT	ND	- weightQuantMode=2/4/5: 1. BS concatenated axis: (T) 2. BS non-concatenated axis: (B*S, 1) - weightQuantMode=3: 1. BS concatenated axis: (T, He/32) 2. BS non-concatenated axis: (B*S, He/32)	×
  dequantScaleWDqOptional	Input	Dequantization parameter of weightDq.	- Supports non-empty tensor (must be passed in weightQuantMode=2/3/4/5 scenarios).	FLOAT	FLOAT8_E8M0, FLOAT	ND	- weightQuantMode=2/4/5: (1, Hcq) - weightQuantMode=3: (Hcq, He/32)	×
  dequantScaleWUqQrOptional	Input	Per-channel parameter for the dequantization operation after MatmulQcQr matrix multiplication.	- Supports non-empty tensor (must be passed in weightQuantMode=1/2/3/4/5 scenarios).	FLOAT	FLOAT, FLOAT8_E8M0	ND	- weightQuantMode=1/2/4/5: (1, N*(D+Dr)) - weightQuantMode=3: (N*(D+Dr), Hcq/32)	×
  dequantScaleWDkvKrOptional	Input	Dequantization parameter of weightDkvKr.	- Supports non-empty tensor (must be passed in weightQuantMode=2/3/4/5 scenarios).	FLOAT	FLOAT8_E8M0, FLOAT	ND	- weightQuantMode=2/4/5: (1, Hckv+Dr) - weightQuantMode=3: (Hckv+Dr, He/32)	×
  quantScaleCkvOptional	Input	Parameter used to quantize the kvCache output data.	- Supports non-empty tensor (must be passed in kvCacheQuantMode=1/2 scenarios).	FLOAT	FLOAT	ND	- kvCacheQuantMode=1: (1) - kvCacheQuantMode=2: (1, Hckv)	×
  quantScaleCkrOptional	Input	Parameter used to quantize the krCache output data.	- Supports non-empty tensor (must be passed in kvCacheQuantMode=2 scenarios).	FLOAT	FLOAT	ND	(1, Dr)	×
  smoothScalesCqOptional	Input	Parameter used for dynamic quantization of the RmsNormCq output.	- Supports non-empty tensor (optional in weightQuantMode=1/2/4/5 scenarios).	FLOAT	FLOAT	ND	(1, Hcq)	×
  actualSeqLenOptional	Input	Indicates the sequence length in each batch, stored in prefix sum form.	- Must be passed when BS concatenated axis and CacheMode="PA_BLK_BSND"/"PA_BLK_NZ".	INT32	INT32	ND	(B)	×
  kNopeClipAlphaOptional	Input	Indicates the scaling factor when performing clip operation on kvCache.	- The shape is 1 in partial quantization and int8 full quantization per-tile scenarios. It can be omitted in other scenarios. Does not support empty tensor.	FLOAT	-	ND	(1)	×
  rmsnormEpsilonCq	Input	The $\epsilon$ parameter in the RmsNorm formula for computing $c^Q$.	- When not explicitly specified by the user, 1e-05 is recommended. - Only the double type is supported.	DOUBLE	DOUBLE	-	-	-
  rmsnormEpsilonCkv	Input	The $\epsilon$ parameter in the RmsNorm formula for computing $c^{KV}$.	- When not explicitly specified by the user, 1e-05 is recommended. - Only the double type is supported.	DOUBLE	DOUBLE	-	-	-
  cacheModeOptional	Input	Indicates the mode of kvCache.	- When not explicitly specified by the user, "PA_BSND" is recommended. - Only the char* type is supported. - Allowed values: "PA_BSND", "PA_NZ", "PA_BLK_BSND", "PA_BLK_NZ", "BSND", "TND".	CHAR*	CHAR*	-	-	-
  queryNormFlag	Input	Indicates whether to output queryNormOutOptional and dequantScaleQNormOutOptional.	- false: do not output; true: output. The default value is false.	BOOL	BOOL	--	--	-
  weightQuantMode	Input	Indicates the quantization mode of weightDq, weightUqQr, weightUk, and weightDkvKr.	- 0: non-quantized; 1: weightUqQr quantized; 2: weightDq, weightUqQr, weightDkvKr int8 quantized; 3: weightDq, weightUqQr, weightDkvKr mxfp8 quantized; 4: weightDq, weightUqQr, weightDkvKr fp8 quantized; 5: weightDq, weightUqQr, weightDkvKr hif8 quantized. The default value is 0.	INT	INT	--	--	-
  kvCacheQuantMode	Input	Indicates the quantization mode of kvCache.	- 0: non-quantized; 1: per-tensor quantization; 2: per-channel quantization; 3: per-tile quantization. The default value is 0.	INT64	INT64	--	--	-
  queryQuantMode	Input	Indicates the quantization mode of Query.	- 0: non-quantized; 1: per-token-head quantization. The default value is 0.	INT64	INT64	--	--	-
  ckvkrRepoMode	Input	Indicates the storage mode of kvCache and krCache.	- 0: kvCache and krCache are stored separately; 1: kvCache and krCache are stored in combined form. The default value is 0.	INT64	INT64	--	--	-
  quantScaleRepoMode	Input	Indicates the storage mode of quantization scale.	- 0: quantization scale and data are stored separately; 1: quantization scale and data are stored in combined form and output as kvCacheRef. The default value is 0.	INT64	INT64	--	--	-
  tileSize	Input	Indicates the size of each tile in per-tile quantization. Valid only when kvCacheQuantMode is 3.	- The default value is 128.	INT64	INT64	--	--	-
  qcQrScale	Input	Indicates the scale correction coefficient of Query.	- When not explicitly specified by the user, 1.0 must be passed.	DOUBLE	DOUBLE	-	-	-
  kcScale	Input	Indicates the scale correction coefficient of Key.	- When not explicitly specified by the user, 1.0 must be passed.	DOUBLE	DOUBLE	-	-	-
  queryOut	Output	Output tensor of Query in the formula (corresponding to $q^N$).	-	BFLOAT16, INT8	BFLOAT16, FLOAT8_E4M3FN, INT8, HIFLOAT8	ND	- BS concatenated axis: (T, N, Hckv) - BS non-concatenated axis: (B, S, N, Hckv)	×
  queryRopeOut	Output	Output tensor of Query positional encoding in the formula (corresponding to $q^R$).	-	BFLOAT16	BFLOAT16	ND	- BS concatenated axis: (T, N, Dr) - BS non-concatenated axis: (B, S, N, Dr)	×
  dequantScaleQNopeOutOptional	Output	Quantization parameter of the Query output in the formula.	- Output when weightQuantMode=2/3/4/5; nullptr when weightQuantMode=0/1.	FLOAT	FLOAT	ND	- BS concatenated axis: (T, N, 1) - BS non-concatenated axis: (B*S, N, 1)	×
  queryNormOutOptional	Output	Output tensor after applying RmsNorm to tokenX in the formula (corresponding to $c^Q$).	- Output when queryNormFlag=true.	BFLOAT16, INT8	BFLOAT16, INT8, FLOAT8_E4M3FN, HIFLOAT8	ND	- BS concatenated axis: (T, Hcq) - BS non-concatenated axis: (B, S, Hcq)	×
  dequantScaleQNormOutOptional	Output	Dequantization parameter of queryNormOutOptional.	- Output when queryNormFlag=true and weightQuantMode=1/2/3/4/5; nullptr when weightQuantMode=0.	FLOAT	FLOAT, FLOAT8_E8M0	ND	- weightQuantMode=1/2/4/5: 1. BS concatenated axis: (T, 1) 2. BS non-concatenated axis: (B*S, 1) - weightQuantMode=3: 1. BS concatenated axis: (T, Hcq/32) 2. BS non-concatenated axis: (B*S, Hcq/32)	×
  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 process.	- 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 required parameter (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 interface (such as the Runtime service not started, memory allocation failure, etc.).
  ACLNN_ERR_INNER_TILING_ERROR	561002	A tiling exception occurs, where the dtype or shape of the input parameter is incorrect.

aclnnMlaPrologV3WeightNz

• Parameters

  Name	Type	Description
  workspace	void*	Device-side allocated workspace memory address.
  workspaceSize	uint64_t	Device-side allocated workspace size, obtained from the first-phase API aclnnMlaPrologV3WeightNzGetWorkspaceSize.
  executor	aclOpExecutor*	Op executor, which contains the operator computation flow.
  stream	aclrtStream	Specifies the stream for task execution.

• Return Value

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

Constraints

• Deterministic computation:

  • aclnnMlaPrologV3WeightNz defaults to a non-deterministic implementation. Deterministic computation can be enabled by calling aclrtCtxSetSysParamOpt.

• Description of shape format fields

  Field Name	Full Name/Meaning	Value Rule and Description
  B	Batch	Value range: 0 to 65536.
  S	Seq-Length	Value range: unlimited.
  He	Head-Size	For A2, A3, and A5, the value is fixed to: 1024, 2048, 3072, 4096, 5120, 6144, 7168, 7680, 8192.
  Hcq	q low-rank matrix dimension	The value is fixed to: 1536.
  N	Head-Num	Value range: 1, 2, 4, 8, 16, 32, 64, 128.
  Hckv	kv low-rank matrix dimension	The value is fixed to: 512.
  D	qk dimension excluding positional encoding	The value is fixed to: 128.
  Dr	qk positional encoding dimension	The value is fixed to: 64.
  Nkv	Number of kv heads	The value is fixed to: 1.
  BlockNum	Number of blocks in the PagedAttention scenario	1. When CacheMode is "PA_BSND"/"PA_NZ", the value is greater than or equal to the result of (B*S)/BlockSize rounded up. 2. When CacheMode is "PA_BLK_BSND"/"PA_BLK_NZ", the value is greater than or equal to the result of B * (S / BlockSize) rounded up (i.e., B * Ceil(S/BlockSize)). Note: In the BS concatenated axis scenario, the S length in each batch can be different, so the value of BlockNum must be greater than or equal to the sum of the rounded-up results of dividing the S length in each batch by BlockSize. Example: actualSeqLenOptional values are [47, 151, 261, 422], blocksize=128, then the lengths in the batch are [47, 104, 110, 161], and BlockNum=Ceil(47/128)+ Ceil(104/128)+ Ceil(110/128)+ Ceil(161/128)=5.
  BlockSize	Block size in the PagedAttention scenario	Value range: 16 to 1024, and must be a multiple of 16.
  T	Size after BS concatenated axis	Value range: unlimited. Note: If BS concatenated axis is used, tokenX, ropeSin, and ropeCos are all 2-dimensional, cacheIndex is 1-dimensional, and queryOut and queryRopeOut are 3-dimensional.
  Dtile	Size of the D dimension of kvCache	- In the per-tile quantization scenario, the value is fixed to 656. - In other scenarios, the value is fixed to Hckv (512).

• Shape constraints

  • If the dimension of tokenX uses BS concatenated axis, i.e., (T, He)

    • The shape of ropeSin and ropeCos is (T, Dr)

    • When CacheMode is PA_BSND or PA_NZ, the shape of cacheIndex is (T)

    • When CacheMode is PA_BLK_BSND or PA_BLK_NZ, the shape of cacheIndex is (Sum(Ceil(S_i/BlockSize))), where S_i is the length of S in each Batch.

    • When CacheMode is PA_BLK_BSND or PA_BLK_NZ, actualSeqLenOptional must be passed in, with the dimension (B).

    • In the int8/fp8/hif8 full quantization scenario, the shape of dequantScaleXOptional is (T, 1); in the mxfp8 full quantization scenario, the shape of dequantScaleXOptional is (T, He/32).

    • The shape of queryOut is (T, N, Hckv).

    • The shape of queryRopeOut is (T, N, Dr).

    • In the int8/mxfp8/fp8/hif8 full quantization scenario, the shape of dequantScaleQNopeOutOptional is (T, N, 1); in other scenarios, it is nullptr.

  • 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).

    • When CacheMode is PA_BSND or PA_NZ, the shape of cacheIndex is (B, S).

    • When CacheMode is PA_BLK_BSND or PA_BLK_NZ, the shape of cacheIndex is (B, Ceil(S/BlockSize)).

    • In the int8/fp8/hif8 full quantization scenario, the shape of dequantScaleXOptional is (B*S, 1); in the mxfp8 full quantization scenario, the shape of dequantScaleXOptional is (B*S, He/32).

    • The shape of queryOut is (B, S, N, Hckv).

    • The shape of queryRopeOut is (B, S, N, Dr).

    • In the int8/mxfp8/fp8/hif8 full quantization scenario, the shape of dequantScaleQNopeOutOptional is (B*S, N, 1); in other scenarios, it is nullptr.

  • One or more of the B, S, T, and Skv values are allowed to be 0, meaning that input parameters whose shapes are related to B, S, T, or Skv values can be empty tensors. Other input parameters do not support empty tensors.

    • If the values of B, S, and T are 0, queryOut and queryRopeOut output empty tensors, and kvCacheRef and krCacheRef are not updated.

    • If the value of Skv is 0, queryOut, queryRopeOut, and dequantScaleQNopeOutOptional are computed normally, while kvCacheRef and krCacheRef are not updated, i.e., they output empty tensors.

  • When CacheMode is BSND:

    • tokenX should not use BS concatenated axis, i.e., its dimensions are (B, S, He).

    • The dimensions of kvCache are (B, S, Nkv, Dr).

  • When CacheMode is TND:

    • tokenX uses BS concatenated axis, with dimensions (T, He).

    • The dimensions of kvCache are (T, Nkv, Dr).

  • When ckvkrRepoMode = 1:

    • The dimensions of krCache must include 0, supporting dimensions (0).

• Special constraints

  • When actualSeqLenOptional is passed in, the last value of actualSeqLenOptional must be consistent with T.

  • In per-tile quantization mode, ckvkrRepoMode and quantScaleRepoMode must both be 1. In other quantization modes and non-quantized scenarios, ckvkrRepoMode and quantScaleRepoMode must both be 0.

  • In per-tile quantization mode, CacheMode only supports PA_BSND, BSND, and TND.

  • When ckvkrRepoMode is 1, krCache must be an empty tensor (that is, the product of its shape dimensions is 0).

In kvcache per-tensor quantization mode, both kvCacheQuantMode and queryQuantMode must be 1. Scenarios supported by the aclnnMlaPrologV3WeightNz API: A2 and A3 currently do not support fp8/hif8/mxfp8 full quantization scenarios, while A5 currently supports all quantization scenarios.

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

  Parameter	Non-quantized scenario	Partial quantization scenario			int8 full quantization scenario			mxfp8 full quantization scenario			fp8 full quantization scenario			hif8 full quantization scenario.
  		kvCache non-quantized.	kvCache per-channel quantization.	kvCache per-tile quantization.	kvCache non-quantized.	kvCache per-tensor quantization	kvCache per-tile quantization	kvCache non-quantized	kvCache per-tensor quantization	kvCache per-tile quantization	kvCache non-quantized	kvCache per-tensor quantization	kvCache per-tile quantization	kvCache non-quantized	kvCache per-tensor quantization	kvCache per-tile quantization
  	dtype	dtype	dtype	dtype	dtype	dtype	dtype	dtype	dtype	dtype	dtype	dtype	dtype	dtype	dtype	dtype
  tokenX	BFLOAT16	BFLOAT16	BFLOAT16	BFLOAT16	INT8	INT8	INT8	FLOAT8_E4M3FN	FLOAT8_E4M3FN	FLOAT8_E4M3FN	FLOAT8_E4M3FN	FLOAT8_E4M3FN	FLOAT8_E4M3FN	HIFLOAT8	HIFLOAT8	HIFLOAT8
  weightDq	BFLOAT16	BFLOAT16	BFLOAT16	BFLOAT16	INT8	INT8	INT8	FLOAT8_E4M3FN	FLOAT8_E4M3FN	FLOAT8_E4M3FN	FLOAT8_E4M3FN	FLOAT8_E4M3FN	FLOAT8_E4M3FN	HIFLOAT8	HIFLOAT8	HIFLOAT8
  weightUqQr	BFLOAT16	INT8	INT8	INT8	INT8	INT8	INT8	FLOAT8_E4M3FN	FLOAT8_E4M3FN	FLOAT8_E4M3FN	FLOAT8_E4M3FN	FLOAT8_E4M3FN	FLOAT8_E4M3FN	HIFLOAT8	HIFLOAT8	HIFLOAT8
  weightUk	BFLOAT16	BFLOAT16	BFLOAT16	BFLOAT16	BFLOAT16	BFLOAT16	BFLOAT16	BFLOAT16	BFLOAT16	BFLOAT16	BFLOAT16	BFLOAT16	BFLOAT16	BFLOAT16	BFLOAT16	BFLOAT16
  weightDkvKr	BFLOAT16	BFLOAT16	BFLOAT16	BFLOAT16	INT8	INT8	INT8	FLOAT8_E4M3FN	FLOAT8_E4M3FN	FLOAT8_E4M3FN	FLOAT8_E4M3FN	FLOAT8_E4M3FN	FLOAT8_E4M3FN	HIFLOAT8	HIFLOAT8	HIFLOAT8
  rmsnormGammaCq	BFLOAT16	BFLOAT16	BFLOAT16	BFLOAT16	BFLOAT16	BFLOAT16	BFLOAT16	BFLOAT16	BFLOAT16	BFLOAT16	BFLOAT16	BFLOAT16	BFLOAT16	BFLOAT16	BFLOAT16	BFLOAT16
  rmsnormGammaCkv	BFLOAT16	BFLOAT16	BFLOAT16	BFLOAT16	BFLOAT16	BFLOAT16	BFLOAT16	BFLOAT16	BFLOAT16	BFLOAT16	BFLOAT16	BFLOAT16	BFLOAT16	BFLOAT16	BFLOAT16	BFLOAT16
  ropeSin	BFLOAT16	BFLOAT16	BFLOAT16	BFLOAT16	BFLOAT16	BFLOAT16	BFLOAT16	BFLOAT16	BFLOAT16	BFLOAT16	BFLOAT16	BFLOAT16	BFLOAT16	BFLOAT16	BFLOAT16	BFLOAT16
  ropeCos	BFLOAT16	BFLOAT16	BFLOAT16	BFLOAT16	BFLOAT16	BFLOAT16	BFLOAT16	BFLOAT16	BFLOAT16	BFLOAT16	BFLOAT16	BFLOAT16	BFLOAT16	BFLOAT16	BFLOAT16	BFLOAT16
  kvCacheRef	BFLOAT16	BFLOAT16	INT8	INT8	BFLOAT16	INT8	INT8	BFLOAT16	FLOAT8_E4M3FN	FLOAT8_E4M3FN	BFLOAT16	FLOAT8_E4M3FN	FLOAT8_E4M3FN	BFLOAT16	HIFLOAT8	HIFLOAT8
  krCacheRef	BFLOAT16	BFLOAT16	INT8	BFLOAT16	BFLOAT16	BFLOAT16	BFLOAT16	BFLOAT16	BFLOAT16	BFLOAT16	BFLOAT16	BFLOAT16	BFLOAT16	BFLOAT16	BFLOAT16	BFLOAT16
  cacheIndexOptional	INT64	INT64	INT64	INT64	INT64	INT64	INT64	INT64	INT64	INT64	INT64	INT64	INT64	INT64	INT64	INT64
  dequantScaleXOptional	NULLPTR	NULLPTR	NULLPTR	NULLPTR	FLOAT	FLOAT	FLOAT	FLOAT8_E8M0	FLOAT8_E8M0	FLOAT8_E8M0	FLOAT	FLOAT	FLOAT	FLOAT	FLOAT	FLOAT
  dequantScaleWDqOptional	NULLPTR	NULLPTR	NULLPTR	NULLPTR	FLOAT	FLOAT	FLOAT	FLOAT8_E8M0	FLOAT8_E8M0	FLOAT8_E8M0	FLOAT	FLOAT	FLOAT	FLOAT	FLOAT	FLOAT
  dequantScaleWUqQrOptional	NULLPTR	FLOAT	FLOAT	FLOAT	FLOAT	FLOAT	FLOAT	FLOAT8_E8M0	FLOAT8_E8M0	FLOAT8_E8M0	FLOAT	FLOAT	FLOAT	FLOAT	FLOAT	FLOAT
  dequantScaleWDkvKrOptional	NULLPTR	NULLPTR	NULLPTR	NULLPTR	FLOAT	FLOAT	FLOAT	FLOAT8_E8M0	FLOAT8_E8M0	FLOAT8_E8M0	FLOAT	FLOAT	FLOAT	FLOAT	FLOAT	FLOAT
  quantScaleCkvOptional	NULLPTR	NULLPTR	FLOAT	NULLPTR	NULLPTR	FLOAT	NULLPTR	NULLPTR	FLOAT	NULLPTR	NULLPTR	FLOAT	NULLPTR	NULLPTR	FLOAT	NULLPTR
  quantScaleCkrOptional	NULLPTR	NULLPTR	FLOAT	NULLPTR	NULLPTR	NULLPTR	NULLPTR	NULLPTR	NULLPTR	NULLPTR	NULLPTR	NULLPTR	NULLPTR	NULLPTR	NULLPTR	NULLPTR
  smoothScalesCqOptional	NULLPTR	FLOAT	FLOAT	FLOAT	FLOAT	FLOAT	FLOAT	NULLPTR	NULLPTR	NULLPTR	FLOAT	FLOAT	FLOAT	FLOAT	FLOAT	FLOAT
  kNopeClipAlphaOptional	NULLPTR	NULLPTR	NULLPTR	FLOAT	NULLPTR	NULLPTR	FLOAT	NULLPTR	NULLPTR	NULLPTR	NULLPTR	NULLPTR	FLOAT	NULLPTR	NULLPTR	FLOAT
  queryOut	BFLOAT16	BFLOAT16	BFLOAT16	BFLOAT16	BFLOAT16	INT8	BFLOAT16	BFLOAT16	FLOAT8_E4M3FN	BFLOAT16	BFLOAT16	FLOAT8_E4M3FN	BFLOAT16	BFLOAT16	HIFLOAT8	BFLOAT16
  queryRopeOut	BFLOAT16	BFLOAT16	BFLOAT16	BFLOAT16	BFLOAT16	BFLOAT16	BFLOAT16	BFLOAT16	BFLOAT16	BFLOAT16	BFLOAT16	BFLOAT16	BFLOAT16	BFLOAT16	BFLOAT16	BFLOAT16
  dequantScaleQNopeOutOptional	NULLPTR	NULLPTR	NULLPTR	NULLPTR	NULLPTR	FLOAT	NULLPTR	NULLPTR	FLOAT	NULLPTR	NULLPTR	FLOAT	NULLPTR	NULLPTR	FLOAT	NULLPTR
  queryNormOutOptional	BFLOAT16	INT8	INT8	INT8	INT8	INT8	INT8	FLOAT8_E4M3FN	FLOAT8_E4M3FN	FLOAT8_E4M3FN	FLOAT8_E4M3FN	FLOAT8_E4M3FN	FLOAT8_E4M3FN	HIFLOAT8	HIFLOAT8	HIFLOAT8
  dequantScaleQNormOutOptional	NULLPTR	FLOAT	FLOAT	FLOAT	FLOAT	FLOAT	FLOAT	FLOAT8_E8M0	FLOAT8_E8M0	FLOAT8_E8M0	FLOAT	FLOAT	FLOAT	FLOAT	FLOAT	FLOAT

Example

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

#include <iostream>
#include <cstring>
#include <vector>
#include <cstdint>
#include "acl/acl.h"
#include "aclnnop/aclnn_mla_prolog_v3_weight_nz.h"
#include<unistd.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, 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-side 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 = aclrtMallocHost(hostAddr, size);
    CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("aclrtMallocHost failed. ERROR: %d\n", ret); return ret);
    memset(*hostAddr, 0, size);
    // 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-side 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 = aclrtMallocHost(hostAddr, size);
    CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("aclrtMallocHost failed. ERROR: %d\n", ret); return ret);
    memset(*hostAddr, 0, size);
    // 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 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) {
    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 for device/stream initialization. See the AscendCL API list.
    // Fill in the deviceId based on your actual device.
    int32_t deviceId = 0;
    aclrtStream stream;
    auto ret = Init(deviceId, &stream);
    // Handle the check based on your needs.
    CHECK_RET(ret == 0, LOG_PRINT("Init acl failed. ERROR: %d\n", ret); return ret);
    // 2. Construct the 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};              // 1
    std::vector<int64_t> smoothScalesCqShape = {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* smoothScalesCqDeviceAddr = 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* smoothScalesCqHostAddr = 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* smoothScalesCq = nullptr;
    int64_t weightQuantMode = 2;
    int64_t kvQuantMode = 1;
    int64_t queryQuantMode = 1;
    int64_t ckvkrRepoMode = 0;
    int64_t quantScaleRepoMode = 0;
    int64_t tileSize = 128;
    double kNopeClipAlpha = 1.0f;
    double qcQrScale = 1.0f;
    double kcScale = 1.0f;
    aclTensor* query = nullptr;
    aclTensor* queryRope = nullptr;
    aclTensor* dequantScaleQNope = 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_INT8_SIZE);
    CHECK_RET(ret == 0, LOG_PRINT("trans NZ shape failed.\n"); return ret);
    ret = TransToNZShape(weightUqQrShape, EXAMPLE_INT8_SIZE);
    CHECK_RET(ret == 0, LOG_PRINT("trans NZ shape failed.\n"); return ret);
    ret = TransToNZShape(weightDkvKrShape, EXAMPLE_INT8_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_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 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 cacheIndex aclTensor.
    ret = CreateAclTensorND(cacheIndexShape, &cacheIndexDeviceAddr, &cacheIndexHostAddr, aclDataType::ACL_INT64, &cacheIndex);
    CHECK_RET(ret == ACL_SUCCESS, return ret);
    // Create kvCache aclTensor.
    ret = CreateAclTensorND(kvCacheShape, &kvCacheDeviceAddr, &kvCacheHostAddr, aclDataType::ACL_INT8, &kvCache);
    CHECK_RET(ret == ACL_SUCCESS, return ret);
    // Create 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 smoothScalesCq aclTensor.
    ret = CreateAclTensorND(smoothScalesCqShape, &smoothScalesCqDeviceAddr, &smoothScalesCqHostAddr, aclDataType::ACL_FLOAT, &smoothScalesCq);
    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 with the specific API.
    uint64_t workspaceSize = 0;
    aclOpExecutor* executor = nullptr;
    // Call the first-phase API of aclnnMlaPrologV3WeightNz.
    ret = aclnnMlaPrologV3WeightNzGetWorkspaceSize(tokenX, weightDq, weightUqQr, weightUk, weightDkvKr, rmsnormGammaCq, rmsnormGammaCkv, ropeSin, ropeCos, kvCache, krCache, cacheIndex,
      dequantScaleX, dequantScaleWDq, dequantScaleWUqQr, dequantScaleWDkvKr, quantScaleCkv, nullptr, smoothScalesCq, nullptr, nullptr,rmsnormEpsilonCq, rmsnormEpsilonCkv, cacheMode,
      weightQuantMode, kvQuantMode, queryQuantMode, ckvkrRepoMode, quantScaleRepoMode, tileSize, qcQrScale, kcScale,
      query, queryRope, dequantScaleQNope, nullptr, nullptr, &workspaceSize, &executor);
    CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("aclnnMlaPrologV3WeightNzGetWorkspaceSize failed. ERROR: %d\n", ret); return ret);
    // Apply for device memory based on the workspaceSize calculated 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 aclnnMlaPrologV3WeightNz.
    ret = aclnnMlaPrologV3WeightNz(workspaceAddr, workspaceSize, executor, stream);
    CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("aclnnMlaPrologV3WeightNz 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 this 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(dequantScaleX);
    aclDestroyTensor(dequantScaleWDq);
    aclDestroyTensor(dequantScaleWUqQr);
    aclDestroyTensor(dequantScaleWDkvKr);
    aclDestroyTensor(quantScaleCkv);
    aclDestroyTensor(smoothScalesCq);
    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(smoothScalesCqDeviceAddr);
    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(smoothScalesCqHostAddr);
    aclrtFree(queryHostAddr);
    aclrtFree(queryRopeHostAddr);
    aclrtFree(dequantScaleQNopeHostAddr);

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

    return 0;
}

The following is an A5 sample code for reference only.

#include <iostream>
#include <cstring>
#include <vector>
#include <cstdint>
#include "acl/acl.h"
#include "aclnnop/aclnn_mla_prolog_v3_weight_nz.h"
#include <unistd.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 = aclrtMallocHost(hostAddr, size);
  CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("aclrtMallocHost failed. ERROR: %d\n", ret); return ret);
  memset(*hostAddr, 0, size);
  // 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 memory.
  ret = aclrtMallocHost(hostAddr, size);
  CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("aclrtMallocHost failed. ERROR: %d\n", ret); return ret);
  memset(*hostAddr, 0, size);
  // 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 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;
}

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.
  // Set deviceId based on your actual device.
  int32_t deviceId = 0;
  aclrtStream stream;
  auto ret = Init(deviceId, &stream);
  // Handle the check based on your 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, 24576};        // Hcq,N*(D+Dr)
  std::vector<int64_t> weightUkShape = {128, 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> kvCacheShape = {1, 16, 1, 512};      // BlockNum,BlockSize,Nkv,Hckv
  std::vector<int64_t> krCacheShape = {1, 16, 1, 64};       // BlockNum,BlockSize,Nkv,Dr
  std::vector<int64_t> cacheIndexShape = {8, 1};              // B,S
  std::vector<int64_t> dequantScaleXShape = {8, 224};           // B*S, 1
  std::vector<int64_t> dequantScaleWDqShape = {1536, 224};      // 1, Hcq
  std::vector<int64_t> dequantScaleWUqQrShape = {24576, 48};    // 1, N*(D+Dr)
  std::vector<int64_t> dequantScaleWDkvKrShape = {576, 224};    // 1, Hckv+Dr
  std::vector<int64_t> quantScaleCkvShape = {1};    // 1
  std::vector<int64_t> queryShape = {8, 1, 128, 512};          // B,S,N,Hckv
  std::vector<int64_t> queryRopeShape = {8, 1, 128, 64};       // B,S,N,Dr
  std::vector<int64_t> dequantScaleQNopeShape = {8, 128, 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* 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* 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* kvCache = nullptr;
  aclTensor* krCache = nullptr;
  aclTensor* cacheIndex = nullptr;
  aclTensor* dequantScaleX = nullptr;
  aclTensor* dequantScaleWDq = nullptr;
  aclTensor* dequantScaleWUqQr = nullptr;
  aclTensor* dequantScaleWDkvKr = nullptr;
  aclTensor* quantScaleCkv = nullptr;
  int64_t weightQuantMode = 3;
  int64_t kvQuantMode = 1;
  int64_t queryQuantMode = 1;
  int64_t ckvkrRepoMode = 0;
  int64_t quantScaleRepoMode = 0;
  int64_t tileSize = 128;
  double qcQrScale = 1.0f;
  double kcScale = 1.0f;
  aclTensor* query = nullptr;
  aclTensor* queryRope = nullptr;
  aclTensor* dequantScaleQNope = 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_INT8_SIZE);
  CHECK_RET(ret == 0, LOG_PRINT("trans NZ shape failed.\n"); return ret);
  ret = TransToNZShape(weightUqQrShape, EXAMPLE_INT8_SIZE);
  CHECK_RET(ret == 0, LOG_PRINT("trans NZ shape failed.\n"); return ret);
  ret = TransToNZShape(weightDkvKrShape, EXAMPLE_INT8_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_FLOAT8_E4M3FN, &tokenX);
  CHECK_RET(ret == ACL_SUCCESS, return ret);
  // Create the weightDq aclTensor.
  ret = CreateAclTensorNZ(weightDqShape, &weightDqDeviceAddr, &weightDqHostAddr, aclDataType::ACL_FLOAT8_E4M3FN, &weightDq);
  CHECK_RET(ret == ACL_SUCCESS, return ret);
  // Create the weightUqQr aclTensor.
  ret = CreateAclTensorNZ(weightUqQrShape, &weightUqQrDeviceAddr, &weightUqQrHostAddr, aclDataType::ACL_FLOAT8_E4M3FN, &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_FLOAT8_E4M3FN, &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_FLOAT8_E4M3FN, &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 dequantScaleX aclTensor.
  ret = CreateAclTensorND(dequantScaleXShape, &dequantScaleXDeviceAddr, &dequantScaleXHostAddr, aclDataType::ACL_FLOAT8_E8M0, &dequantScaleX);
  CHECK_RET(ret == ACL_SUCCESS, return ret);
  // Create dequantScaleWDq aclTensor.
  ret = CreateAclTensorND(dequantScaleWDqShape, &dequantScaleWDqDeviceAddr, &dequantScaleWDqHostAddr, aclDataType::ACL_FLOAT8_E8M0, &dequantScaleWDq);
  CHECK_RET(ret == ACL_SUCCESS, return ret);
  // Create dequantScaleWUqQr aclTensor.
  ret = CreateAclTensorND(dequantScaleWUqQrShape, &dequantScaleWUqQrDeviceAddr, &dequantScaleWUqQrHostAddr, aclDataType::ACL_FLOAT8_E8M0, &dequantScaleWUqQr);
  CHECK_RET(ret == ACL_SUCCESS, return ret);
  // Create dequantScaleWDkvKr aclTensor.
  ret = CreateAclTensorND(dequantScaleWDkvKrShape, &dequantScaleWDkvKrDeviceAddr, &dequantScaleWDkvKrHostAddr, aclDataType::ACL_FLOAT8_E8M0, &dequantScaleWDkvKr);
  CHECK_RET(ret == ACL_SUCCESS, return ret);
  // Create quantScaleCkv aclTensor.
  ret = CreateAclTensorND(quantScaleCkvShape, &quantScaleCkvDeviceAddr, &quantScaleCkvHostAddr, aclDataType::ACL_FLOAT, &quantScaleCkv);
  CHECK_RET(ret == ACL_SUCCESS, return ret);
  // Create the query aclTensor.
  ret = CreateAclTensorND(queryShape, &queryDeviceAddr, &queryHostAddr, aclDataType::ACL_FLOAT8_E4M3FN, &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 aclnnMlaPrologV3WeightNz.
  ret = aclnnMlaPrologV3WeightNzGetWorkspaceSize(tokenX, weightDq, weightUqQr, weightUk, weightDkvKr, rmsnormGammaCq, rmsnormGammaCkv, ropeSin, ropeCos, kvCache, krCache, cacheIndex,
      dequantScaleX, dequantScaleWDq, dequantScaleWUqQr, dequantScaleWDkvKr, quantScaleCkv, nullptr, nullptr, nullptr, nullptr, rmsnormEpsilonCq, rmsnormEpsilonCkv, cacheMode,
      weightQuantMode, kvQuantMode, queryQuantMode, ckvkrRepoMode, quantScaleRepoMode, tileSize, qcQrScale, kcScale,
      query, queryRope, dequantScaleQNope, nullptr, nullptr, &workspaceSize, &executor);
  CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("aclnnMlaPrologV3WeightNzGetWorkspaceSize failed. ERROR: %d\n", ret); return ret);
  // Apply for device memory based on the workspaceSize calculated 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 aclnnMlaPrologV3WeightNz.
  ret = aclnnMlaPrologV3WeightNz(workspaceAddr, workspaceSize, executor, stream);
  CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("aclnnMlaPrologV3WeightNz failed. ERROR: %d\n", ret); return ret);

  // 4. Boilerplate: synchronously 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 side. 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(dequantScaleX);
  aclDestroyTensor(dequantScaleWDq);
  aclDestroyTensor(dequantScaleWUqQr);
  aclDestroyTensor(dequantScaleWDkvKr);
  aclDestroyTensor(quantScaleCkv);
  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(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(queryHostAddr);
  aclrtFree(queryRopeHostAddr);
  aclrtFree(dequantScaleQNopeHostAddr);

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

  _exit(0);
}