FA Supporting the Prefix Encoder

Description

The FA supports the kvCache stored by block. The BlockTables dimension is [batch, maxBlockNum]. Each batch corresponds to the index of the key and value corresponding to the current query. maxBlockNum is approximately equal to the longest kvSeqLen in the key and value. kvSeqLen/blockSize is rounded up.

How to Enable

Set calcType to PREFIX_ENCODER.

Constraints

• If qScale is not supported, set it to 1.
• isTriuMask must be set to 1.
• kernelType must be set to KERNELTYPE_HIGH_PRECISION because float32 is used for KV calculation.
• The clamp is not supported. clampMin and clampMax must be set to 0.
• inputLayout supports only TYPE_BSND.
• If mla is not supported, mlaVHeadSize must be set to 0.
• PrefixCacheLen is calculated by subtracting seqLen from kvseqlen. Due to implementation restrictions, prefixCacheLen must be 128-aligned.
• The restrictions on headNum are as follows:
  • headNum > 0
  • headNum >= kvHeadNum
  • headNum is an integer multiple of kvHeadNum.
• maskType can be MASK_TYPE_ALIBI_COMPRESS, MASK_TYPE_ALIBI_COMPRESS_SQRT, MASK_TYPE_NORM_COMPRESS, or MASK_TYPE_CAUSAL_MASK.