CANN Lightning Indexer Prolog算子文档

custom_pypto.npu_lightning_indexer_prolog_pto【免费下载链接】cann-recipes-infer本项目针对LLM与多模态模型推理业务中的典型模型、加速算法提供基于CANN平台的优化样例项目地址: https://gitcode.com/cann/cann-recipes-infer产品支持情况产品是否支持Atlas A3 系列产品√功能说明算子功能用于 Deepseek IndexerAttention 中计算 Lightning Indexer 所需要的 querykey 和 weights。 Indexer Prolog 的量化策略如下Q_b_proj 使用 W8A8 量化其他 Linear 均不量化query 使用 A8 量化key(cache) 使用 C8 量化反量化因子以 FP16 存储weights 以 FP16 存储Query 的计算公式如下$$ \bold{q}, \bold{q}{scale} \text{DynamicQuant}(\text{Hadamard}(\text{RoPE}(\text{DeQuant}(\bold{q} \cdot \bold{w}{qb})))) $$Q 的计算采用了动态的 Per-Token-Head 量化其中 Hadamard 变换通过矩阵右乘 hadamard_q 实现。而 $\bold{q}, \bold{w}_{qb}$ 均是 Int8 类型。Key(cache) 的计算公式如下$$ \bold{k}, \bold{k}_{scale} \text{DynamicQuant}(\text{Hadamard}(\text{RoPE}(\text{LayerNorm}(\bold{x} \cdot \bold{w}_k)))) $$Cache 的计算同样采用了动态的 Per-Token-Head 量化其中 Hadamard 变换通过矩阵右乘 hadamard_k 实现。Weights 的计算公式如下$$ \bold{weight} (\bold{x} \cdot \bold{w}_{proj}) * \text{scale} $$Weights 的计算没有采用量化同时需要最后转化为 FP16 数据类型供后续的 Lightning Indexer 计算使用。函数原型custom_pypto.npu_lightning_indexer_prolog_pto(token_x, q_norm, q_norm_scale, wq_b, wq_b_scale, wk, weights_proj, ln_gamma_k, ln_beta_k, cos_idx_rope, sin_idx_rope,hadamard_q, hadamard_k, idx_k_cache, idx_k_scale_cache, idx_k_cache_index, layernorm_epsilon_k, layout_queryTND, layout_keyPA_BSND) - (Tensor, Tensor, Tensor)参数说明说明token_xTensor表示 hidden 状态必选参数不支持非连续的Tensor数据格式支持ND数据类型支持bfloat16。q_normTensor表示经过 rmsnorm 后量化的 query必选参数不支持非连续的Tensor数据格式支持ND数据类型支持int8。q_norm_scaleTensor表示 query 的反量化因子必选参数不支持非连续的Tensor数据格式支持ND数据类型支持float32。wq_bTensor表示 query 的权重必选参数不支持非连续的Tensor数据格式支持NZ数据类型支持int8。wq_b_scaleTensor表示 query 的权重反量化因子必选参数不支持非连续的Tensor数据格式支持ND数据类型支持float32。wkTensor表示 key 的权重必选参数不支持非连续的Tensor数据格式支持NZ数据类型支持bfloat16。weights_projTensor表示 weights 的权重必选参数不支持非连续的Tensor数据格式支持NZ数据类型支持bfloat16。ln_gamma_kTensor表示 key 的 layernorm 缩放必选参数不支持非连续的Tensor数据格式支持ND数据类型支持bfloat16。ln_beta_kTensor表示 key 的 layernorm 偏移必选参数不支持非连续的Tensor数据格式支持ND数据类型支持bfloat16。cos_idx_ropeTensor表示用于 RoPE 的 cos不支持非连续的 Tensor数据格式支持 ND数据类型支持bfloat16。sin_idx_ropeTensor表示用于 RoPE 的 sin不支持非连续的 Tensor数据格式支持 ND数据类型支持bfloat16hadamard_qTensor表示用于 query Hadamard 变换的权重矩阵不支持非连续的 Tensor数据格式支持 ND数据类型支持bfloat16。hadamard_kTensor表示用于 key Hadamard 变换的权重矩阵不支持非连续的 Tensor数据格式支持 ND数据类型支持bfloat16。idx_k_cacheTensor表示 key 的缓存必选参数不支持非连续的 Tensor数据格式支持 ND数据类型支持int8。idx_k_scale_cacheTensor表示 key 反量化因子的缓存必选参数不支持非连续的 Tensor数据格式支持 ND数据类型支持float16。idx_k_cache_indexTensor表示更新 key 缓存的位置必选参数不支持非连续的 Tensor数据格式支持 ND数据类型支持int64。layernorm_epsilon_kfloat表示 key layernorm 防除 0 系数必选参数数据类型支持float32。layout_querystr可选参数用于标识输入query的数据排布格式默认值TND。当前仅支持 TND。layout_keystr可选参数用于标识输入key的数据排布格式默认值PA_BSND。当前仅支持 PA_BSND。返回值说明queryTensor公式中 query 的输出 tensor数据格式支持 ND数据类型支持int8。query_scaleTensor公式中 query 反量化因子的输出 tensor数据格式支持 ND数据类型支持float16。weightsTensor公式中 weights 的输出 tensor数据格式支持 ND数据类型支持float16。约束说明该接口支持推理场景下使用。该接口支持图模式PyTorch 2.1版本。该接口与PyTorch配合使用时需要保证CANN相关包与PyTorch相关包的版本匹配。算子代码执行示例算子源码执行参考test_lightning_indexer_prolog.py【免费下载链接】cann-recipes-infer本项目针对LLM与多模态模型推理业务中的典型模型、加速算法提供基于CANN平台的优化样例项目地址: https://gitcode.com/cann/cann-recipes-infer创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考