Fix SparseFlashMla small q_head handling on A5 - #3
Draft
yjyang62 wants to merge 2 commits into
Draft
Conversation
Co-authored-by: yjyang62 <yjyang62@users.noreply.github.com>
Co-authored-by: yjyang62 <yjyang62@users.noreply.github.com>
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
描述
修复 SparseFlashMla 在 A5/arch35 路径下处理较小
q_head(16/32)时可能出现输出异常的问题。主要改动:
q_headgroup)不超过 32 时,不再拆成两个 subBlock 各算一半,而是由 subBlock0 处理完整 M,subBlock1 跳过 0 行任务。q_head=32从原来的 16+16 改为 32+0,避开小半块 vector/VF 处理路径。q_head=16从原来的 8+8 改为 16+0,避免 8 行半块触发异常读写。ProcessVec1对 0 行任务增加早返回并释放同步,避免空任务继续进入 VF API。q_head=32、q_head=16的 CSA prefill 回归用例,以及q_head=16的 SWA prefill 回归用例。关联的Issue
未关联特定 Issue。
测试
python3 -m py_compile attention/sparse_flash_mla/tests/pytest/sparse_flash_mla_paramset.pygit diff --check<= 256U / constInfo.qHeadNum循环模式已不存在。说明:当前环境缺少
torch和torch_npu,无法执行依赖 NPU 的 sparse_flash_mla pytest。文档更新
无。
类型标签