CANN 推理 FA 入参调试工具函数实战基于debug_fa_all_params一键定位 FlashAttention 精度问题【免费下载链接】cann-recipes-infer本项目针对LLM与多模态模型推理业务中的典型模型、加速算法提供基于CANN平台的优化样例项目地址: https://gitcode.com/cann/cann-recipes-infer导读在基于 CANN 平台的 LLM 推理优化中FlashAttentionFA是精度问题的高发区sparse_mode、atten_mask、input_layout、actual_seq_lengths等入参任何一个配置错误都会导致 Prefill/Decode 输出偏差、NaN/Inf 或 cache 错位。本篇文章围绕cann-recipes-infer仓库中model-infer-precision-debug技能提供的FA 入参调试工具函数debug_fa_all_params()及其 6 个子函数完整讲解每个函数的校验逻辑、参数语义与调用方式并结合仓库内真实模型的 FA 调用源码说明如何在 FA 调用前一键完成全部入参校验。读完本文你将掌握一套可直接复制的 FA 入参自检脚本能在精度对齐阶段快速圈定问题参数。一、为什么需要一套 FA 入参调试工具在昇腾 NPU 上FlashAttention 通过torch.ops.npu.npu_fused_infer_attention_score或npu_prompt_flash_attention/npu_incre_flash_attention等算子执行。其入参横跨 layout、mask、长度序列、量化、rope 等多个维度且FA v1 与 v2 的参数名完全不同如 v1 用scale/num_heads/actual_seq_lengthsv2 用softmax_scale/num_query_heads/actual_seq_qlen混用参数名时算子不会报错而是悄悄使用默认值最终表现为精度偏差。从仓库的 model-infer-precision-debug 技能说明 可以确认该技能将“FA 入参检查”列为快速诊断阶段的固定动作而完整的校验实现就集中在本篇要讲解的工具文件中fa_debug_utils.md。使用方式一句话总结将工具代码复制到模型调试文件中在 FA 调用前插入debug_fa_all_params()一键校验全部入参也可以按需单独调用各子函数做针对性检查。二、函数索引总览编号函数用途F.1debug_atten_mask()atten_mask sparse_mode 联合校验F.2debug_fa_version_params()FA v1/v2 参数名混用检查F.3debug_fa_layout_shape()input_layout 与 tensor shape 匹配F.4debug_actual_seq_lengths()actual_seq_lengths 构造校验F.5debug_inner_precise()inner_precise 行无效修正检查F.6debug_mla_rope()MLA rope 参数校验F.7debug_fa_all_params()一键综合校验调用 F.1–F.6其中 F.7 是总入口它内部自动推导batch_size、q_seq_len、kv_seq_len并依次调用 F.3、F.1、F.4、F.5、F.6最后统一打印校验结论。日常调试中建议优先使用 F.7定位到具体嫌疑参数后再用对应子函数深挖。三、完整工具代码将以下代码复制到模型调试文件中在 FA 调用前插入对应的 debug 函数即可import torch def debug_atten_mask(atten_mask, sparse_mode, q_seq_len, kv_seq_len, batch_size, is_prefill, has_rope, layer_idx0): 校验 atten_mask 与 sparse_mode 的配置正确性 print(f--- Layer {layer_idx} atten_mask Debug ---) print(f sparse_mode: {sparse_mode}) print(f is_prefill: {is_prefill}, q_seq_len: {q_seq_len}, kv_seq_len: {kv_seq_len}) if not is_prefill and sparse_mode ! 0: if q_seq_len 1 and not has_rope: print(f [WARN] Decode sparse_mode{sparse_mode}, Q_S1 无 rope 时被忽略 f但建议显式设为 0 避免歧义) else: print(f [ERROR] Decode 阶段 sparse_mode 应为 0, 当前为 {sparse_mode}) if not is_prefill and atten_mask is not None: print(f [WARN] Decode 阶段传入了 atten_mask (shape{atten_mask.shape}) fsparse_mode0 时通常不需要 mask) if is_prefill and sparse_mode not in [2, 3]: print(f [WARN] Prefill 阶段 sparse_mode 通常为 2(leftUpCausal) 或 3(rightDownCausal), f当前为 {sparse_mode}) if atten_mask is None: if sparse_mode in [1, 2, 3, 4]: print(f [ERROR] sparse_mode{sparse_mode} 要求必须传入 atten_mask但收到 None) else: print(f [OK] sparse_mode0 maskNone: 不做 mask 操作) return mask_shape atten_mask.shape print(f mask shape: {mask_shape}, dtype: {atten_mask.dtype}) if atten_mask.dtype not in [torch.bool, torch.int8, torch.uint8]: print(f [ERROR] atten_mask dtype 必须为 bool/int8/uint8, 当前为 {atten_mask.dtype}) if sparse_mode in [2, 3, 4]: valid_shapes [(2048, 2048), (1, 2048, 2048), (1, 1, 2048, 2048)] if mask_shape not in [torch.Size(s) for s in valid_shapes]: print(f [ERROR] sparse_mode{sparse_mode} 要求 mask shape 为 f(2048,2048)/(1,2048,2048)/(1,1,2048,2048), 当前为 {mask_shape}) else: print(f [OK] mask shape 符合 sparse_mode{sparse_mode} 要求) if sparse_mode 3 and mask_shape in [torch.Size(s) for s in valid_shapes]: mask_2d atten_mask.view(2048, 2048) upper_tri torch.triu(torch.ones(2048, 2048, dtypetorch.bool), diagonal1) if atten_mask.dtype torch.bool: if not mask_2d[upper_tri.bool()].all(): print(f [WARN] sparse_mode3 的 mask 上三角部分未全部为 True) elif sparse_mode in [0, 1]: if len(mask_shape) 2: mask_q_s, mask_kv_s mask_shape[-2:] if mask_q_s q_seq_len: print(f [ERROR] mask Q_S 维度 ({mask_q_s}) 小于 query 实际长度 ({q_seq_len})) if mask_kv_s kv_seq_len: print(f [ERROR] mask KV_S 维度 ({mask_kv_s}) 小于 KV 实际长度 ({kv_seq_len})) print(f [REMIND] 若使用 PAmask 最后一维需 block_table 第二维 × block_size) def debug_fa_version_params(fa_call_kwargs, fa_version): 检查 FA 调用参数是否与版本匹配 v1_only_params {actual_seq_lengths, scale, num_heads, antiquant_scale, antiquant_offset, antiquant_mode, quant_scale1, quant_scale2, quant_offset2, dequant_scale1, dequant_scale2, softmax_lse_flag} v2_only_params {actual_seq_qlen, actual_seq_kvlen, softmax_scale, num_query_heads, dequant_scale_key, dequant_scale_value, dequant_offset_key, dequant_offset_value, key_quant_mode, value_quant_mode, quant_scale_out, quant_offset_out, return_softmax_lse, learnable_sink, dequant_scale_query} wrong_version v1_only_params if fa_version v2 else v2_only_params for param in fa_call_kwargs: if param in wrong_version: print(f [ERROR] 参数 {param} 属于 FA {v1 if fa_version v2 else v2} f但当前使用 FA {fa_version}) if fa_version v2: if softmax_scale not in fa_call_kwargs: print(f [WARN] FA v2 未传 softmax_scale将使用默认值 1.0) if num_query_heads not in fa_call_kwargs: print(f [WARN] FA v2 未传 num_query_heads将使用默认值 1) elif fa_version v1: if scale not in fa_call_kwargs: print(f [WARN] FA v1 未传 scale将使用默认值 1.0) if num_heads not in fa_call_kwargs: print(f [WARN] FA v1 未传 num_heads将使用默认值 1) def debug_fa_layout_shape(query, key, value, input_layout, num_query_heads, num_kv_heads, block_tableNone, layer_idx0): 校验 input_layout 与 QKV tensor shape 的匹配性 print(f--- Layer {layer_idx} FA Layout Debug ---) print(f input_layout: {input_layout}) print(f Q: {query.shape}, K: {key.shape}, V: {value.shape}) q_shape query.shape k_shape key.shape layout_base input_layout.split(_)[0] if layout_base BSH: if len(q_shape) ! 3: print(f [ERROR] BSH layout 要求 Q 为 3D, 当前 {len(q_shape)}D) elif layout_base BNSD: if len(q_shape) ! 4: print(f [ERROR] BNSD layout 要求 Q 为 4D, 当前 {len(q_shape)}D) elif q_shape[1] ! num_query_heads: print(f [ERROR] BNSD Q N轴 ({q_shape[1]}) ! num_query_heads ({num_query_heads})) elif layout_base TND: if len(q_shape) ! 3: print(f [ERROR] TND layout 要求 Q 为 3D, 当前 {len(q_shape)}D) elif q_shape[1] ! num_query_heads: print(f [ERROR] TND Q N轴 ({q_shape[1]}) ! num_query_heads ({num_query_heads})) if block_table is not None: print(f [INFO] PA 模式, KV 按 block_table 索引) if block_table.ndim ! 2: print(f [ERROR] block_table 必须为 2D, 当前 {block_table.ndim}D) def debug_actual_seq_lengths(actual_seq_qlen, actual_seq_kvlen, input_layout, is_prefill, batch_size, q_seq_len, kv_seq_len, layer_idx0): 校验 actual_seq_lengths 的构造正确性 layout_base input_layout.split(_)[0] if layout_base TND: if actual_seq_qlen is None: print(f [ERROR] TND layout 必须传入 actual_seq_qlen) return qlen_list actual_seq_qlen if isinstance(actual_seq_qlen, list) else actual_seq_qlen.tolist() for i in range(1, len(qlen_list)): if qlen_list[i] qlen_list[i-1]: print(f [ERROR] TND actual_seq_qlen 必须单调递增 (cumsum), f但 [{i-1}]{qlen_list[i-1]} [{i}]{qlen_list[i]}) elif layout_base in [BSH, BNSD, BSND]: if actual_seq_kvlen is not None: kvlen_list actual_seq_kvlen if isinstance(actual_seq_kvlen, list) else actual_seq_kvlen.tolist() for i, kv_l in enumerate(kvlen_list): if kv_l kv_seq_len: print(f [ERROR] actual_seq_kvlen[{i}]{kv_l} KV S维度 ({kv_seq_len})) def debug_inner_precise(atten_mask, inner_precise, q_seq_len, sparse_mode, layer_idx0): 检查是否需要开启行无效修正 if q_seq_len 1 or sparse_mode not in [0, 1] or atten_mask is None: return if atten_mask.dtype torch.bool: all_masked_rows atten_mask.all(dim-1) else: all_masked_rows (atten_mask ! 0).all(dim-1) num_invalid_rows all_masked_rows.sum().item() if num_invalid_rows 0 and inner_precise in [0, 1]: print(f [WARN] Layer {layer_idx}: {num_invalid_rows} 个全遮蔽行, f建议设置 inner_precise2 或 3) def debug_mla_rope(query, key, query_rope, key_rope, input_layout, sparse_mode, layer_idx0): 校验 MLA rope 参数 if (query_rope is None) ! (key_rope is None): print(f [ERROR] query_rope 和 key_rope 必须同时配置或同时不配置) return if query_rope is None: return q_D query.shape[-1] rope_D query_rope.shape[-1] if q_D not in [512, 128]: print(f [ERROR] MLA query D 仅支持 512/128, 当前 {q_D}) if rope_D ! 64: print(f [ERROR] rope D 必须为 64, 当前 {rope_D}) if q_D 512 and sparse_mode not in [0, 3, 4]: print(f [ERROR] MLA D512 仅支持 sparse_mode 0/3/4, 当前 {sparse_mode}) def debug_fa_all_params(query, key, value, *, fa_versionv1, input_layoutBSH, sparse_mode0, atten_maskNone, actual_seq_qlenNone, actual_seq_kvlenNone, num_heads1, num_kv_heads0, scale1.0, block_tableNone, query_ropeNone, key_ropeNone, inner_precise0, is_prefillTrue, layer_idx0, **kwargs): FA 入参一键综合校验 print(f\n{*60}) print(fFA {fa_version} 综合校验 - Layer {layer_idx} ({Prefill if is_prefill else Decode})) print(f{*60}) batch_size query.shape[0] if input_layout.startswith(B) else None q_seq_len query.shape[1] if input_layout.startswith(B) else query.shape[0] kv_seq_len key.shape[1] if len(key.shape) 3 else key.shape[0] if scale 1.0: print(f [WARN] scale 1.0, 确认是否忘记设置 1/sqrt(head_dim)?) has_rope query_rope is not None debug_fa_layout_shape(query, key, value, input_layout, num_heads, num_kv_heads, block_table, layer_idx) debug_atten_mask(atten_mask, sparse_mode, q_seq_len, kv_seq_len, batch_size or 1, is_prefill, has_rope, layer_idx) debug_actual_seq_lengths(actual_seq_qlen, actual_seq_kvlen, input_layout, is_prefill, batch_size or 1, q_seq_len, kv_seq_len, layer_idx) debug_inner_precise(atten_mask, inner_precise, q_seq_len, sparse_mode, layer_idx) debug_mla_rope(query, key, query_rope, key_rope, input_layout, sparse_mode, layer_idx) print(f{*60}\n)四、子函数逐项解读F.1–F.6F.1debug_atten_maskatten_mask 与 sparse_mode 联合校验这是整套工具中规则最密集的函数其校验逻辑与技能文档附录 F.1 的 sparse_mode 语义表一一对应sparse_mode含义atten_mask 要求适用路径0defaultMask可选通常NoneMLA absorb Decodeq_len11allMask必传(Q_S, KV_S)特殊场景2leftUpCausal不推荐—3rightDownCausal必传(2048, 2048)下三角标准 LLM TND PAPrefillDecode 统一、MLA absorb Prefill4band滑窗必传(2048, 2048)滑窗层 PrefillDecode 统一函数按四条主线输出诊断Decode 阶段 sparse_mode 检查Decodeis_prefillFalse时sparse_mode应为 0但存在一个例外分支——当q_seq_len 1且has_ropeFalse时算子内部会忽略该 sparse_mode因此只打 WARN 并建议显式设为 0 避免歧义。Decode 误传 atten_masksparse_mode0时通常不需要 mask传入了则提示 WARN典型错误见技能文档 5.4MLA absorb Decode 切到sparse_mode0后没把 mask 设None。mask 缺失检查sparse_mode ∈ {1,2,3,4}时 mask 为None会直接报 ERROR。mask 合法性检查dtype 必须为bool/int8/uint8对sparse_mode ∈ {2,3,4}校验 shape 必须是(2048,2048)/(1,2048,2048)/(1,1,2048,2048)对sparse_mode3额外用torch.triu构造上三角布尔矩阵检查 mask 的上三角部分是否全部为 TruerightDownCausal 要求下三角 mask对sparse_mode ∈ {0,1}校验 mask 后两维不小于实际 Q/KV 长度。函数最后固定输出一条 REMIND若使用 PA分页注意力mask 最后一维需 ≥block_table第二维 ×block_size这对应技能文档快速验证清单第 13 条。F.2debug_fa_version_paramsFA v1/v2 参数名混用检查FA v1 与 v2 的接口差异是精度问题的经典来源。函数内置两份参数名集合v1 专属actual_seq_lengths、scale、num_heads、antiquant_scale/offset/mode、quant_scale1/scale2/offset2、dequant_scale1/scale2、softmax_lse_flagv2 专属actual_seq_qlen、actual_seq_kvlen、softmax_scale、num_query_heads、dequant_scale_key/value、dequant_offset_key/value、key_quant_mode、value_quant_mode、quant_scale_out、quant_offset_out、return_softmax_lse、learnable_sink、dequant_scale_query只要在调用 kwargs 中发现属于另一版本的参数名立即报 ERROR。此外还会检查必需参数是否缺失v2 未传softmax_scale→ 默认 1.0未传num_query_heads→ 默认 1v1 未传scale→ 默认 1.0未传num_heads→ 默认 1。这一点尤其致命参数名传错时算子不报错只是静默使用默认值 1既不是1/sqrt(head_dim)也不是真实 head 数输出精度立刻崩溃。技能文档附录 F.2 给出的 v1/v2 参数映射可与此函数对照使用。F.3debug_fa_layout_shapeinput_layout 与 tensor shape 匹配函数先取input_layout.split(_)[0]得到 layout 基底如TND_NTD→TND再按基底分派校验layoutQ 维度适用场景TND3D[T, N, D]框架部署 PA 标准 LLM 默认推荐T packed tokensBSH3D[B, S, H*D]migrator 骨架 / 独立部署连续缓存非 PABNSD4D[B, N, S, D]非 PA扩散模型对BNSD与TND还会进一步核对 Q 的 N 轴是否等于num_query_heads。若传入block_table则打印 PA 模式提示并校验其必须为 2D。PA 模式下 KV 按 block_table 索引物理存储 shape 为[total_blocks, block_size, num_kv_heads, head_dim]FA 调用时需 view 为[bn, bs, num_head*dim]见技能文档附录 F.3。F.4debug_actual_seq_lengthsactual_seq_lengths 构造校验actual_seq_lengths 的构造方式必须与 layout 匹配layoutactual_seq_qlenactual_seq_kvlenTND累积和cumsum([s1, s2, ...])累积和BSH / BNSD原始值[s1, s2, ...]原始值函数对TND路径校验actual_seq_qlen不能为 None且必须单调递增cumsum 性质否则报 ERROR对BSH/BNSD/BSND路径校验actual_seq_kvlen每一项不得大于 KV 的 S 维度。Decode 阶段actual_seq_qlen每个值通常为 1MTP 场景为next_n1。这一点在仓库框架部署中有直接对应实现在 execution_engine.py 中Prefill 通过actual_seq_lengths_kv seq_lens、actual_seq_lengths_cu_q actual_seq_lengths_cu_kv.cumsum(0)构造累积和Decode 阶段则基于kv_len 1再取 cumsum——与 F.4 的规则完全吻合。F.5debug_inner_precise行无效修正检查函数统计 mask 中被整行遮蔽全遮蔽行的数量对 bool mask 用atten_mask.all(dim-1)对非 bool mask 用(atten_mask ! 0).all(dim-1)。当存在全遮蔽行且inner_precise ∈ {0, 1}时提示需要开启行无效修正值精度模式行无效修正适用场景0高精度否默认1高性能否精度要求不高2高精度是Prefill 自定义 mask 有全遮蔽行3高性能是同上但允许精度损失注意该函数在q_seq_len 1典型 Decode或sparse_mode ∉ {0,1}或 mask 为 None 时直接跳过——行无效修正只对 Prefill 自定义 mask 场景有意义。F.6debug_mla_ropeMLA rope 参数校验针对 MLAMulti-head Latent Attention路径的 rope 约束query_rope与key_rope必须同时传或同时不传rope D 必须为64query D 仅支持512 或 128query D512 时仅支持sparse_mode ∈ {0, 3, 4}。这组约束在仓库 DeepSeek 系列模型的 MLA 实现中有实证在 modeling_deepseek.py 的 FA 调用中query_rope/key_rope成对传入input_layoutNTD_TND、sparse_mode3、atten_maskattention_mask、scaleself.softmax_scale同时给出且qk_rope_head_dim与 64 对齐。五、总入口debug_fa_all_params一键综合校验F.7debug_fa_all_params以关键字参数方式聚合了上述全部子函数的输入注意*之后全部为 keyword-only防止位置错乱。调用前它会自动从 tensor shape 推导batch_sizeinput_layout以B开头时取query.shape[0]否则为 Noneq_seq_leninput_layout以B开头取query.shape[1]否则取query.shape[0]kv_seq_lenkey维度 ≥3 时取key.shape[1]否则取key.shape[0]。随后依次执行scale 默认值告警 → F.3 layout 校验 → F.1 mask 校验 → F.4 长度序列校验 → F.5 行无效检查 → F.6 MLA rope 校验最后打印分隔线收尾。推荐的接入方式# 在 FA 调用前插入示例参数按实际调用填充 debug_fa_all_params( query, key, value, fa_versionv1, input_layoutNTD_TND, sparse_mode3, atten_maskattention_mask, actual_seq_qlenactual_seq_lengths_cu_q, actual_seq_kvlenactual_seq_lengths_kv, num_headsself.num_heads_per_rank, num_kv_headsself.num_heads_per_rank, scaleself.softmax_scale, is_prefillis_prefill, )一个完整的仓库实证DeepSeek R1 的 MLA FA 调用上面接入示例并非虚构它取自仓库 models/deepseek_r1/models/modeling_deepseek.py 的真实 FA 调用形态。该文件在forward_attention_absorb与forward_page_attention_absorb中分别处理 Prefill/Decode 两条路径FA 调用关键参数如下Prefill 路径input_layoutNTD_TND、sparse_mode3、atten_maskattention_mask、actual_seq_lengthsactual_seq_lengths_kv、scaleself.softmax_scale、antiquant_mode0Decode 路径forward_page_attention_absorb根据is_prefill分支切换sparse_mode 3 if is_prefill else 0见该文件约 1184–1186 行与 1296–1298 行对应技能文档所述“标准 LLM TND PA 统一 sparse_mode3MLA absorb Decode 切 sparse_mode0 maskNone”的路径分类。在将该模型从基线 eager 实现迁移到 FA 路径后即可用debug_fa_all_params对着这份调用核对NTD_TND的 Q 3D 性、N 轴与num_heads_per_rank的一致性、cumsum 形式的actual_seq_lengths等。算子侧底层实现可进一步参考 ops/cannbot_dsl/flash_kda.pyFused Chunk Kimi Delta Attention 的 CANN Bot DSL 内核与ops/ascendc/torch_ops_extension/custom_ops下npu_quant_sparse_flash_mla、npu_kv_quant_sparse_attn_sharedkv等算子的注册与转换实现。六、结合调试工作流何时用哪个函数参照 model-infer-precision-debug 技能 的“症状分类 → 快速验证 → 分模块定位”流程可以将本工具函数与典型症状对应起来精度症状优先调用的调试函数典型根因Prefill 正确、Decode 偏差大F.7 / F.1Decode 误用 sparse_mode3 或误传 maskMLA absorb 路径actual_seq_qlen 构造错误输出含 NaN/InfF.7scale 告警、F.2scale 传错名或漏传、dtype 不匹配、量化参数不匹配首个 token 正确、后续偏移F.4Decode kv_len 更新错误、cumsum 构造错误多 batch 部分样本错误F.3、F.1block_table 非 2D、mask 维度小于实际长度量化模式偏差放大F.2v1 误用dequant_*或 v2 误用antiquant_*全遮蔽行导致的输出异常F.5Prefill 自定义 mask 有全遮蔽行但inner_precise仍为 0/1MLA 模型输出错乱F.6query_rope/key_rope 只传其一、rope D≠64、D512 用了不支持的 sparse_mode调试纪律技能文档“重要原则”所有精度判定必须基于与基线eager 模式 / 优化前的数值对比一次只改一个变量保留中间 tensor 与日志实际执行测试。本工具函数只负责“入参自检”输出[ERROR]/[WARN]/[OK]/[REMIND]四类结论其中[ERROR]确定的非法配置必须先修复再跑 FA[WARN]不一定会出错但属于高嫌疑项或易产生歧义的写法[OK]该维度校验通过[REMIND]PA 场景的约束提醒如 mask 最后一维与 block_table 的关系。修复后建议结合技能文档第四步的“逐层对比框架”做最终验证在 Block 输入、Attention Norm 后、QKV 投影后、RoPE 后、KVCache 写入后、FA 输出、Attention 残差后、Block 输出共 8 个位置插入 checkpoint用max_abs_diff / max_rel_diff / cosine_similarity对比基线BF16 模式max_rel 1e-3量化模式max_rel 1e-2确认首个偏差位置确实被消除。七、总结fa_debug_utils.md提供的这 7 个函数F.1–F.7覆盖了 FlashAttention 入参的全部高风险维度sparse_mode 与 atten_mask 的联合约束、FA v1/v2 参数名隔离、input_layout 与 tensor shape 匹配、actual_seq_lengths 构造规则、inner_precise 行无效修正、MLA rope 约束。它们的价值在于把“精度问题要靠猜”变成“入参问题用打印定位”——在 FA 调用前插一行debug_fa_all_params(...)几十条规则一次跑完配合仓库中 DeepSeek 系列模型如 modeling_deepseek.py的真实调用与框架侧 execution_engine.py 的长度序列构造实现即可快速锁定问题参数避免在 NaN/Inf 与 token 错位中盲目排查。【免费下载链接】cann-recipes-infer本项目针对LLM与多模态模型推理业务中的典型模型、加速算法提供基于CANN平台的优化样例项目地址: https://gitcode.com/cann/cann-recipes-infer创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考