FA 调用完整代码示例【免费下载链接】cannbot-skillsCANNBot 是面向 CANN 开发的用于提升开发效率的系列智能体本仓库为其提供可复用的 Skills 模块。项目地址: https://gitcode.com/cann/cannbot-skills基于仓库中已有模型的实际调用按模式分类。模式一连续缓存GPT-OSSTND layout, FA v2, sliding window# 参考: cann-recipes-infer/models/gpt_oss/models/modeling_gpt_oss.py attn_output, _ torch_npu.npu_fused_infer_attention_score_v2( query_states, past_key, past_value, num_query_headsself.num_attention_heads_per_rank, num_key_value_headsself.num_key_value_heads_per_rank, input_layoutTND, softmax_scaleself.scaling, sparse_mode4 if self.sliding_window else 3, pre_tokensself.sliding_window if self.sliding_window else torch.iinfo(torch.int32).max, next_tokens0, actual_seq_qlenactual_seq_qlen, actual_seq_kvlenactual_seq_lengths_kv, atten_maskattention_mask, learnable_sinkself.sinks, )Qwen3-MoEBSH layout, FA v1, Prefill/Decode 分离# 参考: cann-recipes-infer/models/qwen3_moe/models/modeling_qwen3_moe.py # Decode attn_output, _ torch.ops.npu.npu_fused_infer_attention_score( query_states, past_key_states, past_value_states, num_headsself.num_heads_per_rank, num_key_value_headsself.num_key_value_heads_per_rank, input_layoutBSH, atten_maskattention_mask, scaleself.scale_fa, actual_seq_lengths_kvactual_seq_lengths_kv, ) # Prefill注sparse_mode2 为仓库历史实现推荐使用 sparse_mode3 attn_output, _ torch.ops.npu.npu_fused_infer_attention_score( query_states, key_states, value_states, num_headsself.num_heads_per_rank, num_key_value_headsself.num_key_value_heads_per_rank, input_layoutBSH, atten_maskattention_mask, sparse_mode2, scaleself.scale_fa, next_tokens0, )模式二 模式三分页注意力 MLA 压缩缓存以下示例均使用 PAblock_table MLAkeyvaluecache_nope, query_rope/key_rope 分离。DeepSeek-R1TND_NTD layout, FA v2, MLA absorb# 参考: cann-recipes-infer/models/deepseek_r1/models/modeling_deepseek.py attn_output, _ self.fa_ops.npu_fused_infer_attention_score_v2( q_nope, k_nope, k_nope, query_ropeq_pe, key_ropek_rope, atten_maskattention_mask, actual_seq_kvlenactual_seq_lengths_kv, actual_seq_qlenactual_seq_lengths_q, block_tableself.block_table, num_query_headsself.num_heads_per_rank, num_key_value_headsself.num_key_value_heads_per_rank, softmax_scaleself.softmax_scale, input_layoutTND_NTD, sparse_mode0, block_sizeself.block_size, query_quant_mode0, key_quant_mode0, value_quant_mode0, )Kimi-K2TND_NTD, FA v1, Prefill/Decode 分离实例# 参考: cann-recipes-infer/models/kimi-k2-thinking/models/modeling_deepseek.py fa_input_kwargs { query: q_nope, key: k_nope, value: k_nope, query_rope: q_pe, key_rope: k_pe, num_heads: self.num_heads_per_rank, num_key_value_heads: self.num_key_value_heads_per_rank, input_layout: TND_NTD, actual_seq_lengths: actual_seq_qlen, actual_seq_lengths_kv: actual_seq_lengths_kv, sparse_mode: 3, atten_mask: attention_mask, block_table: block_table, block_size: self.block_size, scale: self.softmax_scale, } if is_prefill: attn_output, _ self.fa_ops_prefill.npu_fused_infer_attention_score(**fa_input_kwargs) else: attn_output, _ self.fa_ops_decode.npu_fused_infer_attention_score(**fa_input_kwargs)LongCat-FlashBSND_NBSD, FA v1, KVP# 参考: cann-recipes-infer/models/longcat-flash/models/modeling_longcat_flash.py attn_partial, lse_partial self.fa_ops.npu_fused_infer_attention_score( query_states[0], k_nope, k_nope, query_ropequery_states[1], key_ropek_rope, num_headsself.num_heads_per_rank, num_key_value_headsself.num_key_value_heads_per_rank, input_layoutBSND_NBSD, block_tableself.block_table, block_sizeself.block_size, atten_maskattention_mask, actual_seq_lengths_kvactual_seq_lengths_kv, scaleself.softmax_scale, sparse_modesparse_mode, softmax_lse_flagself.kvp_size 1, )缓存写入融合算子# 参考: cann-recipes-infer/models/longcat-flash/models/modeling_longcat_flash.py _, _, k_rope, k_nope torch_npu.npu_kv_rmsnorm_rope_cache( latent_cache, self.kv_a_layernorm.weight, cos, sin, slot_mapping.view(-1), # 写入位置 rope_cache, nope_cache, # 输出缓存 epsilon1e-6, cache_modePA_NZ, is_output_kvTrue, )【免费下载链接】cannbot-skillsCANNBot 是面向 CANN 开发的用于提升开发效率的系列智能体本仓库为其提供可复用的 Skills 模块。项目地址: https://gitcode.com/cann/cannbot-skills创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考