HySparse MQA sparse attention (FA4 block scoring + SWA MQA + block-sparse DSA gather)

[ForFishes]

What

Add HySparse MQA sparse attention (NSA-style, three cooperating structures) to PaddleFleet, plus test coverage.

Core (b9f3145)

• FA4-fused block scoring (tilelang_ops/hysparse/block_score_fa4.py) with PyLayer-routed gradients.
• SWA sliding-window MQA (tilelang_ops/hysparse/swa_attn.py), replacing the standalone block-sparse MQA fwd/bwd kernels.
• Block-sparse DSA gather (cudnn_ops/block_sparse_mqa_dsa.py) via cuDNN / FlashMLA sparse, including:
  • d_v < 512 value-region padding to 512 so the FlashMLA sparse kernel constraint (d_v == 512) is satisfied transparently — zeros injected between the value region and RoPE region are transparent to q·k, output sliced back to kv_lora_rank; pad/slice sit outside the PyLayer so autograd routes dq/dkv correctly.
  • learnable per-head attn_sink (analytic gradient in backward).
• MLA / transformer / gpt_layer_specs wiring; enable_hy_sparse_attention config plumbing in transformer_config.
• Runtime guards (independent of develop behaviour):
  • fusion_layer_utils: deepep_topk_to_sonic_metadata_with_scales ImportError fallback so older paddlefleet_ops binaries (pre-Supports Sonicmoe bf16 z and wgrad, optimize sonicmoe performance #1348) still load.
  • transformer_config: when enable_hy_sparse_attention, the MTP portion of window_attn_skip_freq must be all-zero (SWA MTP layers have no shared KV) — fail fast.

Tests (da832e0)

FA4 block-score grad / topk-consistency / multidoc, MQA gather DSA, windowed SWA, pipeline, whole-flow precision, MQA self-attention module test, and build_hysparse_valid_range test. Obsolete test_hysparse_block_attn.py removed (superseded).

FlashMLA sparse kernel constraints

d_qk ∈ {512, 576}, d_v == 512, h_q == 64. Absorbed-MLA MQA layout: Dk = kv_lora_rank + qk_rope_head_dim, Dv = kv_lora_rank.

Compatibility

Rebased on latest develop. No changes to existing develop features — HySparse is gated behind enable_hy_sparse_attention. Net diff: 28 files, +4975 / −1979.

Test plan

• Single-card HySparse tests added under tests/single_card_tests/.
• End-to-end 8-card pretraining launches green: step 1 loss ≈ 12.297, no NaN/Inf, forward-backward completing.

[mergify]

⚠️ The sha of the head commit of this PR conflicts with #1453. Mergify cannot evaluate rules on this PR. Once #1453 is merged or closed, Mergify will resume processing this PR. ⚠️

[ForFishes]

Duplicate of #1453 (same code, same head branch name). Closing in favor of #1453 which carries the full review history. #1453 head already points at the cleaned-up commits.

[Copilot]

Pull request overview

该 PR 在 PaddleFleet 中引入 HySparse MQA 稀疏注意力（包含 FA4 block scoring + SWA MQA + cuDNN/FlashMLA DSA gather 三个协作结构），并补充了覆盖算子精度、梯度、packed-multidoc 语义与端到端 whole-flow 的单卡测试，用于在 enable_hy_sparse_attention 配置开关下启用新路径、保持对现有 develop 功能的隔离。

Changes:

• 新增 FA4 融合的 block-score（输出 scaled block_logit）与 host 侧 TopK block 选择流水线，并在 MLA 路径中接入（HySparse full layer 产出 shared KV + block indices）。
• 新增 SWA windowed MQA TileLang kernel（支持 Dk≠Dv、可选 attn_sink）及其 PyLayer 封装，并接入 HySparse SWA 层主分支。
• 新增 DSA（FlashMLA sparse fwd + cuDNN DSA bwd）block-sparse gather 后端与 padding/sink 逻辑，并增加多组单卡测试覆盖。

PR 标题与描述检查

• 标题目前不符合仓库要求的 [CLASS]Title 格式；建议按类别补前缀，例如：[Feature] HySparse MQA sparse attention (FA4 + SWA + DSA gather) 或按项目约定选择更合适的 CLASS。
• 描述信息较完整（What/核心改动/约束/兼容性/测试计划均有），符合最低要求。

Reviewed changes

Copilot reviewed 28 out of 28 changed files in this pull request and generated 1 comment.

Show a summary per file

File	Description
src/paddlefleet/transformer/transformer_layer.py	引入 HySparseTransformerLayer 以支持跨层 KV sharing，并在线程/重计算路径中维护 shared KV 的传递。
src/paddlefleet/transformer/transformer_config.py	增加 HySparse 配置开关与 block/topk 参数，并在 __post_init__ 增加 MTP + SWA 的 fail-fast 约束。
src/paddlefleet/models/gpt/gpt_layer_specs.py	HySparse 开关下切换 attention/layer 实现并对不兼容组合显式报错。
src/paddlefleet/tilelang_ops/hysparse/windowed_mqa_attn*.py	windowed MQA kernel fwd/bwd（Dk≠Dv、attn_sink、ragged seqlen bwd guard 等）。
src/paddlefleet/tilelang_ops/hysparse/swa_attn.py	sliding-window MQA attention 的 PyLayer 封装与默认 sm_scale 逻辑。
src/paddlefleet/tilelang_ops/hysparse/block_score_fa4.py	FA4 融合 block scoring 的 PyLayer 封装与 scaled block_logit 语义。
src/paddlefleet/tilelang_ops/hysparse/pipeline.py	TopK block 选择与 block score 恢复（exp(block_logit - lse)）及 TopK 前的 detach 防护。
src/paddlefleet/cudnn_ops/block_sparse_mqa_dsa.py	DSA gather 后端：block->token 展开、sinkless/learnable sink、kv_lora_rank<512 padding、analytic sink grad。
src/paddlefleet/transformer/moe/fusion_layer_utils.py	为旧版 paddlefleet_ops 增加 deepep_topk_to_sonic_metadata_with_scales 的 ImportError fallback。
tests/single_card_tests/**	新增/调整多组 HySparse 单卡测试，覆盖 grad、一致性、multidoc、pipeline、whole-flow precision 等。

[Paddle-CI-Bot]

PaddleFleet Log Analysis

Run #29429655732 · Attempt 1

日志分析报告

流水线名称	问题标签	修复建议	日志片段
Coverage Upload And Check	增量覆盖率不达标（exit 9, CC_FAIL_ON_ERROR: true）	对 block_sparse_mqa_dsa.py、block_score_fa4.py、transformer_layer.py、multi_latent_attention.py 等核心改动文件补充单测，将增量覆盖率从 37% 提升至 ≥90%	报错代码

失败的测试case:

Coverage Upload And Check (job 87400908780)
diff-cover 增量覆盖率检查：37%（要求 ≥90%）
Total: 633 lines，Missing: 393 lines
exit code 9（CC_FAIL_ON_ERROR: true）

根本原因分析:
PR #1470（HySparse MQA sparse attention）新增/修改了 2195 行代码，涉及 14 个 src/paddlefleet/ 文件，但现有单测对其中核心文件的行覆盖严重不足：block_score_fa4.py 0%、transformer_layer.py 6.4%、block_sparse_mqa_dsa.py 13%、multi_latent_attention.py 34.9%，综合增量覆盖仅 37%，未达 90% 阈值。

修复建议:

1. src/paddlefleet/cudnn_ops/block_sparse_mqa_dsa.py（13%）：补充覆盖 lines 68-388 范围内缺失行的单测，重点覆盖 MQA DSA gather 的主分支逻辑。
2. src/paddlefleet/tilelang_ops/hysparse/block_score_fa4.py（0%）：该文件完全未被覆盖，需新增 FA4 block-score 的基础功能单测。
3. src/paddlefleet/transformer/transformer_layer.py（6.4%）：补充 HySparse MQA 相关 transformer layer 路径（lines 1432–1744）的单测。
4. src/paddlefleet/transformer/multi_latent_attention.py（34.9%）：覆盖 MQA forward/backward 中 lines 642–2013 缺失段。
5. src/paddlefleet/transformer/moe/fusion_layer_utils.py（0%）：补充 lines 42-51 的单测。
6. 已有测试文件（test_hysparse_*.py）统计时被标记为 Skipping non-paddlefleet file，说明其 coverage 数据未被纳入，需确认测试用例运行时正确 import 并执行到 src/paddlefleet/ 路径下的代码，或在 coverage 配置中正确 map 源码路径。

---

🔍 准确性记录：请点击评论底部 😊 图标，选择 👍（准确）或 👎（有误），将自动记录到 CI 监控系统

_{🔄 每次 Re-run 后自动更新}