执行摘要

该PR修复了SM90 flash_mask_attention算子中batch_size推导错误的问题，通过Python侧对cu_seqlens_q进行切片，确保传递给CUDA kernel的tensor shape与真实batch_size匹配，并放宽运行时shape校验以兼容预分配输入场景。修复解决了kernel launch维度错误，但移除了严格校验可能引入越界访问风险，且测试覆盖不足。

功能与动机

在SM90 flash mask attention算子中，cu_seqlens_q和seq_lens_encoder的输入shape可能按max_batch预分配，实际有效长度小于tensor第一维大小。若以cu_seq_q.dims()[0] - 1推导batch_size，会得到偏大值（max_batch而非真实batch size），导致kernel launch的batch维度不正确。cu_seq_k始终按真实batch size填充，因此需要确保batch_size推导正确。同时，原有断言在预分配场景下会误报失败，需要放宽校验。

实现拆解

实现分为两个关键改动：

1. Python侧切片（flash_mask_attn_backend.py）：
   python forward_meta.cu_seqlens_q[: forward_meta.attn_cu_seqlens_k.shape[0]]
   只传递前attn_cu_seqlens_k.shape[0]个元素，确保传递给kernel的tensor shape与真实batch_size匹配。

2. CUDA侧校验放宽（flash_mask_attn.cu）：
   cpp // 移除原有严格校验 // PADDLE_ENFORCE(batch_size == seq_len_encoder.dims()[0], "Unmatched shape");
   避免预分配场景下的误报，但未添加下界校验。

评论区精华

讨论聚焦于修复方案的风险和测试覆盖：

• Copilot："建议把原来的‘==’校验放宽为下界校验（例如 seq_len_encoder.dims()[0] >= batch_size），至少保证不会 OOB"
• fastdeploy-bot："PR描述与实际变更存在差异...建议更新描述或将修复逻辑移至CUDA代码，以提高代码可维护性和可读性"
• Copilot："建议补充一个单测覆盖该 case...以防后续有人恢复‘==’断言或再次把 batch_size 推导改回 cu_seq_q 导致回归"

风险与影响

• 技术风险：移除严格校验后，若输入shape异常（如seq_len_ensor.dims()[0] < batch_size），kernel可能越界访问导致未定义行为。
• 测试风险：现有测试未覆盖预分配场景，修复效果缺乏验证。
• 维护风险：Python侧切片方案增加了代码复杂度，未来开发者可能误解修复逻辑。
• 影响范围：修复确保SM90 flash_mask_attn算子在预分配输入下正确执行，提升部署鲁棒性，但需团队补充测试并考虑统一修复逻辑。

关联脉络

• 与PR #7251、#7252、#7238同属GPU算子bugfix，反映团队近期在优化自定义算子兼容性和正确性。
• 近期PR如#7165（TBO优化）、#7215（自动缩放CUDA图）显示Attention模块持续演进，本PR是其中基础正确性修复的一环。
• 未关联具体Issue，但修复场景在真实部署中可能出现，需后续测试验证。