FA 入参调试工具函数

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

在 FA 调用前插入 debug_fa_all_params() 一键校验全部入参，也可单独调用各子函数进行针对性检查。

函数索引

完整代码

使用时将以下代码复制到模型调试文件中，在 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_idx=0):
    """校验 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_S=1 无 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})，"
              f"sparse_mode=0 时通常不需要 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_mode=0 + mask=None: 不做 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, dtype=torch.bool), diagonal=1)
            if atten_mask.dtype == torch.bool:
                if not mask_2d[upper_tri.bool()].all():
                    print(f"  [WARN] sparse_mode=3 的 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] 若使用 PA，mask 最后一维需 >= 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_table=None, layer_idx=0):
    """校验 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_idx=0):
    """校验 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_idx=0):
    """检查是否需要开启行无效修正"""
    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_precise=2 或 3")


def debug_mla_rope(query, key, query_rope, key_rope, input_layout, sparse_mode, layer_idx=0):
    """校验 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 D=512 仅支持 sparse_mode 0/3/4, 当前 {sparse_mode}")


def debug_fa_all_params(query, key, value, *, fa_version="v1", input_layout="BSH",
                        sparse_mode=0, atten_mask=None, actual_seq_qlen=None,
                        actual_seq_kvlen=None, num_heads=1, num_kv_heads=0, scale=1.0,
                        block_table=None, query_rope=None, key_rope=None,
                        inner_precise=0, is_prefill=True, layer_idx=0, **kwargs):
    """FA 入参一键综合校验"""
    print(f"\n{'='*60}")
    print(f"FA {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")

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