在 kernels/automode/a5/flash_atten/scripts/validate_buffer_usage.py 的 compute_ub_usage_nd 和 compute_ub_usage_dn 函数中:
kernels/automode/a5/flash_atten/scripts/validate_buffer_usage.py
compute_ub_usage_nd
compute_ub_usage_dn
k_tile_factor = tile_s1 // cube_s1 vec_s0 = cube_s0 // VEC_CORES // k_tile_factor
当 tile_s1 < cube_s1 时,tile_s1 // cube_s1 结果为 0,后续 cube_s0 // VEC_CORES // k_tile_factor 会触发 ZeroDivisionError。
tile_s1 < cube_s1
tile_s1 // cube_s1
cube_s0 // VEC_CORES // k_tile_factor
ZeroDivisionError
虽然正常运行场景下 tile_s1 >= cube_s1 通常成立,但该脚本作为 buffer 校验工具,应该对异常参数输入做防御性处理,而不是直接崩溃。用户在调整 FlashAttention 的 tiling 参数时,完全可能传入 tile_s1 < cube_s1 的组合来试探 buffer 上限。
tile_s1 >= cube_s1
建议在除法前增加校验:
k_tile_factor = tile_s1 // cube_s1 if k_tile_factor == 0: return {"total_bytes": 0, "overflow": False, "detail": "tile_s1 < cube_s1, skipped"} vec_s0 = cube_s0 // VEC_CORES // k_tile_factor
或者在 main 函数解析参数后做前置校验,提前报错给出清晰提示。
main
谢谢你的建议,修正中
在
kernels/automode/a5/flash_atten/scripts/validate_buffer_usage.py的compute_ub_usage_nd和compute_ub_usage_dn函数中:当
tile_s1 < cube_s1时,tile_s1 // cube_s1结果为 0,后续cube_s0 // VEC_CORES // k_tile_factor会触发ZeroDivisionError。虽然正常运行场景下
tile_s1 >= cube_s1通常成立,但该脚本作为 buffer 校验工具,应该对异常参数输入做防御性处理,而不是直接崩溃。用户在调整 FlashAttention 的 tiling 参数时,完全可能传入tile_s1 < cube_s1的组合来试探 buffer 上限。建议在除法前增加校验:
k_tile_factor = tile_s1 // cube_s1 if k_tile_factor == 0: return {"total_bytes": 0, "overflow": False, "detail": "tile_s1 < cube_s1, skipped"} vec_s0 = cube_s0 // VEC_CORES // k_tile_factor或者在
main函数解析参数后做前置校验,提前报错给出清晰提示。