[object Object]

[object Object][object Object]undefined[object Object]

• Description: Splits the tail axis of the input tensor (kv) into the left half for rms_norm computation and the right half for RoPE computation, and then scatters the computation results to two caches.

• Formulas:

  (1) interleaveRope:

  $$x=kv[...,Dv:]$$ $$x1=x[...,::2]$$ $$x2=x[...,1::2]$$ $$x\_part1=torch.cat((x1,x2),dim=-1)$$ $$x\_part2=torch.cat((-x2,x1),dim=-1)$$ $$y=x\_part1*cos+x\_part2*sin$$

  (2) rmsNorm:

  $$x=kv[...,:Dv]$$ $$square\_x=x*x$$ $$mean\_square\_x=square\_x.mean(dim=-1,keepdim=True)$$ $$rms=torch.sqrt(mean\_square\_x+epsilon)$$ $$y=(x/rms)*gamma$$

[object Object]

Each operator has calls. First, [object Object] is called to obtain the input parameters and compute the required workspace size based on the process. Then, [object Object] is called to perform computation.

[object Object]

• Parameters

  [object Object]

• Returns:

  [object Object]: status code. For details, see .

  The first-phase API implements input parameter validation. The following error codes may be returned.

  [object Object]

[object Object]

• Parameters

  [object Object]

• Returns:

  [object Object]: status code. For details, see .

[object Object]

• The shape-related variables used in the parameter descriptions are defined as follows:
  • Bkv indicates the batch size of the input KV, and Skv indicates the sequence length of the input KV. The size is determined by the user input scenario and is not limited.
  • N indicates the number of heads of the input KV. This operator is closely related to the DeepSeekV3 network structure and supports only the scenario where N is 1.
  • D indicates the head dimension of the input KV. The data Dv required for rms_norm computation and the data Dk required for RoPE computation are obtained by splitting the input KV data D. Therefore, the size of Dk and Dv must meet the following requirement: Dk + Dv = D. In addition, Dk must meet the rope rules. According to the rope rules, Dk must be an even number. If cacheModeOptional is set to NZ (cacheModeOptional is PA_NZ or PA_BLK_NZ), Dk and Dv must be 32-byte aligned.
  • If cacheModeOptional is set to PA (cacheModeOptional is PA, PA_BNSD, PA_NZ, PA_BLK_BNSD, or PA_BLK_NZ), BlockSize must be 32-byte aligned.
  • The alignment value in the preceding 32-byte alignment scenarios is determined by the data type of the cache. Take BlockSize as an example. If the data type of the cache is INT8, HIFLOAT8, FLOAT8E5M2 or FLOAT8E4M3FN, BlockSize%32 must be equal to 0. If the data type of the cache is float16 or bfloat16, BlockSize%16 must be equal to 0. If the dtype of kCacheRef is different from that of ckvCacheRef, BlockSize must meet both BlockSize%32 = 0 and BlockSize%16 = 0.
  • Bcache is the batch size of the input cache, and Scache is the sequence length of the input cache. The size is determined by the user input scenario and is not limited.
  • BlockNum is the number of memory blocks written to the cache. The size is determined by the user input scenario and is not limited.
• Restrictions on index:
  • When cacheModeOptional is set to Norm, the shape is two-dimensional [Bkv, Skv], and the value of index must be in the range of [–1, Scache). The value can be repeated under different Bkv.
  • When cacheModeOptional is set to PA_BNSD or PA_NZ, the shape is one-dimensional [Bkv Skv], and the value of index must be in the range of [–1, BlockNum BlockSize). The value cannot be repeated.
  • When cacheModeOptional is set to PA_BLK_BNSD or PA_BLK_NZ, the shape is one-dimensional [Bkv ceil_div(Skv, BlockSize)], and the value of index must be in the range of [–1, BlockNum BlockSize). The value/BlockSize cannot be repeated.
• Restrictions on quantization scenarios:
  • Supported quantization scenario 1: The data type of kCacheRef is FLOAT16 or BFLOAT16, and the data type of ckvCacheRef is INT8, HIFLOAT8, FLOAT8E5M2, or FLOAT8E4M3FN.
  • Supported quantization scenario 2: The data type of ckvCacheRef is FLOAT16 or BFLOAT16, and the data type of kCacheRef is INT8, HIFLOAT8, FLOAT8E5M2, or FLOAT8E4M3FN.
  • Supported quantization scenario 3: The data types of kCacheRef and ckvCacheRef are the same, which are INT8, HIFLOAT8, FLOAT8E5M2, or FLOAT8E4M3FN.

[object Object]

The following example is for reference only. For details, see .