CANN/cannbot-skills FA调用完整代码示例

FA 调用完整代码示例

【免费下载链接】cannbot-skillsCANNBot 是面向 CANN 开发的用于提升开发效率的系列智能体，本仓库为其提供可复用的 Skills 模块。项目地址: https://gitcode.com/cann/cannbot-skills

基于仓库中已有模型的实际调用，按模式分类。

模式一：连续缓存

GPT-OSS（TND 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_heads=self.num_attention_heads_per_rank, num_key_value_heads=self.num_key_value_heads_per_rank, input_layout="TND", softmax_scale=self.scaling, sparse_mode=4 if self.sliding_window else 3, pre_tokens=self.sliding_window if self.sliding_window else torch.iinfo(torch.int32).max, next_tokens=0, actual_seq_qlen=actual_seq_qlen, actual_seq_kvlen=actual_seq_lengths_kv, atten_mask=attention_mask, learnable_sink=self.sinks, )

Qwen3-MoE（BSH 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_heads=self.num_heads_per_rank, num_key_value_heads=self.num_key_value_heads_per_rank, input_layout="BSH", atten_mask=attention_mask, scale=self.scale_fa, actual_seq_lengths_kv=actual_seq_lengths_kv, ) # Prefill（注：sparse_mode=2 为仓库历史实现，推荐使用 sparse_mode=3） attn_output, _ = torch.ops.npu.npu_fused_infer_attention_score( query_states, key_states, value_states, num_heads=self.num_heads_per_rank, num_key_value_heads=self.num_key_value_heads_per_rank, input_layout="BSH", atten_mask=attention_mask, sparse_mode=2, scale=self.scale_fa, next_tokens=0, )

模式二 + 模式三：分页注意力 + MLA 压缩缓存

以下示例均使用 PA（block_table）+ MLA（key=value=cache_nope, query_rope/key_rope 分离）。

DeepSeek-R1（TND_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_rope=q_pe, key_rope=k_rope, atten_mask=attention_mask, actual_seq_kvlen=actual_seq_lengths_kv, actual_seq_qlen=actual_seq_lengths_q, block_table=self.block_table, num_query_heads=self.num_heads_per_rank, num_key_value_heads=self.num_key_value_heads_per_rank, softmax_scale=self.softmax_scale, input_layout="TND_NTD", sparse_mode=0, block_size=self.block_size, query_quant_mode=0, key_quant_mode=0, value_quant_mode=0, )

Kimi-K2（TND_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-Flash（BSND_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_rope=query_states[1], key_rope=k_rope, num_heads=self.num_heads_per_rank, num_key_value_heads=self.num_key_value_heads_per_rank, input_layout="BSND_NBSD", block_table=self.block_table, block_size=self.block_size, atten_mask=attention_mask, actual_seq_lengths_kv=actual_seq_lengths_kv, scale=self.softmax_scale, sparse_mode=sparse_mode, softmax_lse_flag=self.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, # 输出缓存 epsilon=1e-6, cache_mode="PA_NZ", is_output_kv=True, )

【免费下载链接】cannbot-skillsCANNBot 是面向 CANN 开发的用于提升开发效率的系列智能体，本仓库为其提供可复用的 Skills 模块。项目地址: https://gitcode.com/cann/cannbot-skills