【合入来源】

• 需求 https://gitcode.com/Ascend/op-plugin/issues/102
• 问题单
• issue/工单
• 重构优化
• 资料更新

【修改方案】

1. 适配 npu_block_sparse_attention_backward TND 场景：

   • 在 BlockSparseAttentionBackwardKernelNpuOpApi.cpp 中补充 TND 格式必要参数校验。
   • 当 q_input_layout == "TND" 时，校验 actual_seq_lengths 必传且非空。
   • 当 kv_input_layout == "TND" 时，校验 actual_seq_lengths_kv 必传且非空。
   • 不新增 layout 白名单、layout 组合校验、GQA 头数校验，相关能力边界继续由底层 CANN aclnn 接口判定。

2. 适配并验证 TND/GQA 反向计算：

   • 新增 TND + GQA CPU golden，包括 block sparse mask、TND offsets、GQA head 映射和 KV 梯度聚合逻辑。
   • 新增显式 backward 用例，覆盖 full mask、sparse mask、单 batch、变长、group size 变化、MQA、非默认 block shape 尾块等场景。
   • 新增 autograd 端到端用例，覆盖 npu_block_sparse_attention(...).backward() 路径，并与 CPU golden 对比。
   • 新增 TND 缺失 actual_seq_lengths / actual_seq_lengths_kv 的本地校验负例。

3. 更新文档及文档生成模板：

   • 更新 npu_block_sparse_attention 文档中的反向能力描述。
   • 明确反向路径支持 BNSD/TND 同 layout 场景，并支持 MHA/GQA。
   • 明确 headDim、GQA head 数关系、actual_seq_lengths 与 actual_seq_lengths_kv 配置约束。
   • 同步更新 _op_plugin_docs.py，避免后续文档生成覆盖手工修改。

【资料变更】

涉及资料变更。

修改文件：

• docs/zh/custom_APIs/torch_npu/torch_npu-npu_block_sparse_attention.md
• codegen/templates/_op_plugin_docs.py

主要内容：

• 修正反向能力过期描述。
• 补充 BNSD/TND、MHA/GQA、headDim、actual_seq_lengths / actual_seq_lengths_kv 的约束说明。

【接口变更】

不涉及 public API schema 变更。

说明：

• 未修改 op_plugin_functions.yaml。
• 未新增 npu_block_sparse_attention_backward 对外公开文档。
• 未变更 Python API 入参、出参和调用方式。
• 本次仅补充 TND 必要参数本地校验，其他 layout/GQA 支持范围仍由底层 CANN aclnn 接口判定。

【功能验证】

新增/适配 UT：

• test_npu_block_sparse_attention_backward_tnd_gqa_full_mask_cpu_compare
• test_npu_block_sparse_attention_backward_tnd_gqa_sparse_mask_cpu_compare
• test_npu_block_sparse_attention_backward_tnd_gqa_single_batch_cpu_compare
• test_npu_block_sparse_attention_backward_tnd_gqa_uneven_seq_lengths_cpu_compare
• test_npu_block_sparse_attention_backward_tnd_gqa_group_size_4_cpu_compare
• test_npu_block_sparse_attention_backward_tnd_mqa_cpu_compare
• test_npu_block_sparse_attention_backward_tnd_gqa_non_128_tail_block_cpu_compare
• test_npu_block_sparse_attention_backward_tnd_actual_seq_lengths_required
• test_npu_block_sparse_attention_backward_tnd_gqa_autograd_cpu_compare
  
  
  
  验证方式：
• 执行 python test/test_custom_ops/test_npu_block_sparse_attention_backward.py。
• 显式 backward 用例对比 dq/dk/dv 与 CPU golden。
• autograd 端到端用例对比 query.grad/key.grad/value.grad 与 CPU golden。
• TND 缺失 seq 参数用例验证本地 TORCH_CHECK 报错。

【CheckList】

• 代码注释完备，正确记录错误日志
• 代码实现进行了返回值、空指针等校验
• PR标题正确使用类型标签，如：feat、fix、refactor、docs、test等
• PR持续集成流水线（CI）执行通过，代码检查无异常