Skip to content

Fix SparseFlashMla small q_head handling on A5 - #3

Draft
yjyang62 wants to merge 2 commits into
mainfrom
cursor/fix-sparse-flash-mla-small-heads-9fc3
Draft

Fix SparseFlashMla small q_head handling on A5#3
yjyang62 wants to merge 2 commits into
mainfrom
cursor/fix-sparse-flash-mla-small-heads-9fc3

Conversation

@yjyang62

@yjyang62 yjyang62 commented Jul 7, 2026

Copy link
Copy Markdown
Owner

描述

修复 SparseFlashMla 在 A5/arch35 路径下处理较小 q_head(16/32)时可能出现输出异常的问题。

主要改动:

  • A5 arch35:当单任务 M 维(对应当前 q_head group)不超过 32 时,不再拆成两个 subBlock 各算一半,而是由 subBlock0 处理完整 M,subBlock1 跳过 0 行任务。
    • q_head=32 从原来的 16+16 改为 32+0,避开小半块 vector/VF 处理路径。
    • q_head=16 从原来的 8+8 改为 16+0,避免 8 行半块触发异常读写。
  • A5 arch35:ProcessVec1 对 0 行任务增加早返回并释放同步,避免空任务继续进入 VF API。
  • A2/A3 arch22:顺手修正 SWA/CSA vector block 中 sinks brcb 扩展循环的边界,避免多做一段无效 DMA 写入。
  • 测试参数补充:增加 q_head=32q_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.py
  • git diff --check
  • 使用 ripgrep 确认旧的 <= 256U / constInfo.qHeadNum 循环模式已不存在。

说明:当前环境缺少 torchtorch_npu,无法执行依赖 NPU 的 sparse_flash_mla pytest。

文档更新

无。

类型标签

  • 🐛 Bug修复
  • ✨ 新特性
  • ⚡ 性能优化
  • ♻️ 重构
  • 🧪 测试
  • 📦 构建/CI
  • 🔧 配置变更
  • 📝 文档更新
  • ⬆️ 依赖升级
  • 🔒 安全修复
  • 🧹 代码清理
  • ❓ 其他,请描述:
Open in Web Open in Cursor 

cursoragent and others added 2 commits July 7, 2026 07:11
Co-authored-by: yjyang62 <yjyang62@users.noreply.github.com>
Co-authored-by: yjyang62 <yjyang62@users.noreply.github.com>
@cursor cursor Bot changed the title Fix SparseFlashMla small q_head sinks expansion Fix SparseFlashMla small q_head handling on A5 Jul 7, 2026
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants