MultiLatentAttentionOperation

功能

MLA场景下的paged attention，使用分页管理的kvcache计算attention score，额外支持分离qnope/qrope、ctkv/krope的输入。

定义

struct MultiLatentAttentionParam {
    int32_t headNum = 0;
    float qkScale = 1.0;
    int32_t kvHeadNum = 0;
    enum MaskType : int {
        UNDEFINED = 0,      
        MASK_TYPE_SPEC,      
        MASK_TYPE_MASK_FREE, 
    };
    MaskType maskType = UNDEFINED;
    enum CalcType : int {
        CALC_TYPE_UNDEFINED = 0,
        CALC_TYPE_SPEC,          
        CALC_TYPE_RING,         
    };
    CalcType calcType = CALC_TYPE_UNDEFINED;
    enum CacheMode : uint8_t {
        KVCACHE = 0,  
        KROPE_CTKV,   
        INT8_NZCACHE, 
        NZCACHE,     
    };
    CacheMode cacheMode = KVCACHE;
    uint8_t rsv[43] = {0};
};

参数列表

上表中类型为自定义类型的，其描述如下：

• maskType：表示mask类型，其具体取值如下。
  • UNDEFINED：无mask。
  • MASK_TYPE_SPEC：并行解码mask，与calcType = CALC_TYPE_SPEC一起使用。
  • MASK_TYPE_MASK_FREE：传入固定shape的mask，与calcType = CALC_TYPE_SPEC一起使用。计算时根据尾部长度从传入mask中获取真实的mask。mask为倒三角形式，第一行全部为-inf，最后一行全部为0。具体形状如下：

    def generate_mask_free(q_len):
            :param q_len: q_len
            :return: constructed mask, e.g. when q_len=2, returned as
            [[-10000.0 -10000.0 -10000.0 ... -10000.0],
             [0        -10000.0 -10000.0 ... -10000.0],
             [0               0 -10000.0 ... -10000.0],
             ...
             [0               0        0 ... -10000.0],
             [0               0        0 ...        0]]
            """
            mask_free = np.full((125 + 2 * q_len, 128), -10000.0)
            mask_free = np.triu(mask_free, 2 - q_len)
            return mask_free

• calcType：表示计算类型，其具体取值如下。
  • CALC_TYPE_UNDEFINED：默认的decoder场景。
  • CALC_TYPE_SPEC：支持qseqlen大于1，并行解码功能，MTP使用该场景。
  • CALC_TYPE_RING：ringAttention，额外输出Ise。
• cacheMode：表示输入query和kcache的类型，其具体取值如下。

  与MLAPO大融合预处理算子对应使用。

  • KVCACHE：输入只有一个q和一个合并kvCache，为原PagedAttention下的MLA场景，目前暂未接入。
  • KROPE_CTKV：输入的q拆分为qNope和qRope，输入的kcache拆分为ctKV和kRope，对应之前的默认场景。
  • INT8_NZCACHE：高性能cache，在KROPE_CTKV的基础上：krope和ctkv转为NZ格式输出，ctkv和qnope经过per_head静态对称量化为int8类型。
  • NZCACHE：在KROPE_CTKV的基础上krope和ctkv转为NZ格式输出。

输入

输出

规格约束

• block_size ＜= 128，建议为128。
• batch <= 8192
• cacheMode为INT8_NZCACHE时，不支持num_heads = 128