[Feature] Add deterministic TileLang backend for absorbed-MQA sparse backward - #1776
Conversation
There was a problem hiding this comment.
Pull request overview
该 PR 为 absorbed-MQA latent sparse attention 增加可选的 TileLang 确定性后向实现，并通过配置项把后向 kernel 从 cuDNN（快速但 dkv 非 bitwise 可复现）切换为 TileLang（确定性但更慢），同时补齐覆盖 dispatch / 配置透传 / 关键边界形状的单测。
Changes:
- 在
TransformerConfig中新增mqa_sparse_attn_backward_backend配置项，并添加校验与测试覆盖。 - 在
mqa_sparse_attnPyLayer 中新增backward_backend分支，接入tilelang_ops.attn.mqa_latent_sparse_bwd作为 absorbed-MQA 的可选后向实现。 - 新增与扩展测试：覆盖 determinism、与 fp32 oracle 对齐、head-group/chunk/topk padding 等 host 侧 tiling 边界。
PR 标题/描述建议（按仓库约定）
- 标题当前不符合
[CLASS]Title格式；建议改为类似：[Feature] Add deterministic TileLang backend for absorbed-MQA sparse backward（或与实际分类一致的[NewFeature]...）。 - 描述建议补充：引入该开关的动机（复现性/aadiff）、如何启用（配置名/参数名）、以及性能权衡（约 ~14x 慢）和适用场景（仅在需要确定性时启用）。
Reviewed changes
Copilot reviewed 7 out of 7 changed files in this pull request and generated 5 comments.
Show a summary per file
| File | Description |
|---|---|
| tests/single_card_tests/transformer/test_mqa_latent_sparse_bwd.py | 新增 TileLang absorbed-MQA 后向的数值/确定性/host tiling 边界测试 |
| tests/single_card_tests/transformer/test_mqa_latent_attention.py | 增加配置透传与 determinism 的集成测试用例 |
| tests/single_card_tests/transformer/test_dsv4_hybrid_attention.py | 增加新配置项的校验与默认值测试 |
| src/paddlefleet/transformer/transformer_config.py | 新增配置项 mqa_sparse_attn_backward_backend、映射与校验逻辑 |
| src/paddlefleet/transformer/mqa_latent_attention.py | 读取新配置并透传到 sparse attention 后向选择；修复 indexer backend 透传 |
| src/paddlefleet/tilelang_ops/attn/mqa_latent_sparse_bwd.py | 新增 absorbed-MQA 的 TileLang 确定性后向 wrapper（复用 symmetric kernel） |
| src/paddlefleet/fusions/mqa_sparse_attn.py | 后向分支支持 tilelang，并在 forward 记录/校验 backend |
💡 Add a code-review agent skill or configure MCP servers for context-aware, tailored reviews. Learn more in the docs.
| * "cudnn" (default): cuDNN DSA backward. Fast but ``dkv`` accumulates | ||
| with atomics and is **not** run-to-run reproducible (rel ~2e-3). | ||
| * "tilelang": deterministic backward via | ||
| ``tilelang_ops.attn.mqa_latent_sparse_bwd``. ~14x slower on SM100 | ||
| but bitwise stable under ``FLAGS_cudnn_deterministic=1``. |
There was a problem hiding this comment.
已按建议改掉：不再写死 rel ~2e-3，也不再把 tilelang 的 bitwise 稳定性绑定到 FLAGS_cudnn_deterministic。现在文档指向 test_block_sparse_dsa_gradcheck.py::TestDeterminism 作为 cuDNN dkv 漂移的唯一界定处，并说明 tilelang 侧是无条件走 atomic-free 的 bwd_det，与该 flag 无关。
| backend = str(getattr(config, "csa_indexer_backend", "tilelang")) | ||
| self.indexer_backend = "tilelang" if backend == "unfused" else backend | ||
| # Backward kernel for the sparse MQA attention (dkv). cuDNN accumulates | ||
| # dkv with atomics and is not run-to-run reproducible (rel ~2e-3); |
There was a problem hiding this comment.
已改：这里同样去掉了写死的 rel ~2e-3，改为指向 TestDeterminism，并注明 tilelang 的稳定性来自实现（by construction）而不是某个 flag。
| alone; only the backward changes, because the cuDNN DSA backward accumulates | ||
| ``dkv`` with atomics and is therefore not run-to-run reproducible (rel ~2e-3, | ||
| measured -- see the comment at ``mqa_sparse_attn.py:257``), which is the one | ||
| remaining source of step-to-step aadiff in the non-absorbed MQA layers once the | ||
| 38 CSA/HCA layers are on ``csa_sparse_attn_backend=tilelang``. |
There was a problem hiding this comment.
已改：删掉了 mqa_sparse_attn.py:257 这个会失效的行号引用，改为引用 test_block_sparse_dsa_gradcheck.py::TestDeterminism；同时在模块 docstring 的 Determinism 段落补了一句，明确 bwd_det 是无条件调用的，与 FLAGS_cudnn_deterministic 无关（对称版的 sparse_mqa_bwd.py:613 才读该 flag 选 kernel）。
| b, s, h, dk = query.shape | ||
| _, s_kv, kv_dk = kv.shape | ||
| dv = grad_out.shape[-1] | ||
| if block_h is None: | ||
| block_h = _pick_block_h(h) | ||
|
|
There was a problem hiding this comment.
已加显式 dtype 校验：query / kv / out / grad_out 任一不是 bfloat16 就抛 ValueError 并点名该张量（lse / attn_sink 按契约是 fp32，不在校验内）。确认了复用的 kernel 在 JIT 侧固定 bf16（sparse_mqa_bwd.py:33/86/129/337 的 assert dtype == T.bfloat16），fp16 进来会落到 bf16 kernel 上读 fp16 显存。配套单测 test_non_bf16_inputs_are_rejected 逐个张量验证报错信息。
| - ``"cudnn"`` (default): ``paddlefleet.cudnn_ops.csa_sparse_attn_bwd_cudnn`` | ||
| (cuDNN DSA). Fast but ``dkv`` is not run-to-run reproducible (atomics). | ||
| - ``"tilelang"``: ``paddlefleet.tilelang_ops.attn.mqa_latent_sparse_bwd`` | ||
| (deterministic). ~14x slower on SM100 but bitwise stable across runs when | ||
| ``FLAGS_cudnn_deterministic=1``. |
There was a problem hiding this comment.
已改：这段模块说明不再把 bitwise 稳定性写成需要 FLAGS_cudnn_deterministic=1，而是说明它对相同输入恒定稳定、且与该 flag 无关（原因是这里无条件走 atomic-free kernel，而对称版 sparse_mqa_bwd 才按 flag 二选一）；cuDNN 侧的漂移量级改为指向 TestDeterminism。
There was a problem hiding this comment.
| 序号 | 位置 | 优先级 | 状态 |
|---|---|---|---|
| 1 | transformer_config.py:1348 |  |
✅ |
| 2 | DeepGEMM 子模块说明 | |
🟡 |
更正：经 live base 219e6845 与当前 head 的 Git tree 复核，两者的 DeepGEMM 指针均为 60bd61e1，且 GitHub changed-files/compare 均不包含该路径。此前结论来自陈旧的本地 develop 引用，该 P1 撤回，无需本 PR 修改子模块。
| kernel. | ||
| """ | ||
|
|
||
| mqa_sparse_attn_backward_backend: str = "cudnn" |
There was a problem hiding this comment.
config_adapter/precision.py 的精度适配规则。当前 --test-accuracy 会将现有 CSA sparse/indexer backend 切到 tilelang，但该字段仍保留默认 cudnn；配置包含 absorbed-MQA 层时，dKV 反向仍走非确定性实现，精度对齐流程无法获得本 PR 提供的确定性结果。请同时补充对应的 config adapter 测试。
LiYuRio
changed the title
ebc5737 to
abd0c56
Compare
abd0c56 to
e68e7a4
Compare
e68e7a4 to
bf4543b
Compare
risemeup1111
left a comment
risemeup1111
left a comment
There was a problem hiding this comment.
71677099 回退到落后 19 个提交的 60bd61e1，重新引入 SM100 FP8 1d1d bias epilogue 越界风险。请恢复至少 71677099，或使用包含该修复的更新提交。
|
感谢 review。逐项回应： Review Board 第 2 项（DeepGEMM 子模块）—— 这条是误报，本 PR 未改动该子模块
该 OOB 修复确实值得同步，但把第三方子模块 bump 混进这个特性 PR 会让改动范围失焦，建议单独提一个 PR 升级指针（我可以来做）。麻烦复核后放行这一项。 Copilot 的 5 条，均已修
risemeup1111 第 1 项（
|
risemeup1111
left a comment
risemeup1111
left a comment
There was a problem hiding this comment.
71677099 回退到落后 19 个提交的 60bd61e1，重新引入 SM100 FP8 1d1d bias epilogue 越界风险。请恢复至少 71677099，或使用包含该修复的更新提交。
Codecov Report❌ Patch coverage is Additional details and impacted files@@ Coverage Diff @@
## develop #1776 +/- ##
==========================================
Coverage ? 93.90%
==========================================
Files ? 4
Lines ? 164
Branches ? 35
==========================================
Hits ? 154
Misses ? 8
Partials ? 2
Flags with carried forward coverage won't be shown. Click here to find out more.
🚀 New features to boost your workflow:
|
|
❌ Cherry-pick failed: Conflicts detected when cherry-picking to |
PaddlePaddle#1776 turned ``_MQASparseAttention.backward`` into a two-branch dispatch on ``mqa_sparse_attn_backward_backend`` (cuDNN vs the new deterministic tilelang kernel), re-indenting the whole cuDNN body under an ``else:``. This branch had edited two lines inside that body, so the wholesale move and the small edit could not be reconciled automatically. Resolution keeps develop's structure verbatim and re-threads ``global_kv_idx_remap_fusion`` through it: - forward takes both new keywords; the fused remap still feeds the FlashMLA forward, which runs regardless of the backward backend. - the local->global KV column remap moves into the ``"cudnn"`` branch only. The tilelang backward indexes ``token_indices`` per batch itself and never builds a flat-global table, so the switch has no meaning there. - ``mqa_sparse_attn`` passes ``global_kv_idx_remap_fusion`` before ``backward_backend``, matching the forward's parameter order. The two fusion switches in ``_FakeCtx`` and the ``mqa_sparse_attn`` call in ``MQALatentAttention._sparse_attn`` were both-added conflicts: keep both. ``test_mqa_sparse_attn_end_to_end`` now sweeps the backward backend, since the combination is new. It also lets the test pin ``dkv``, which the cuDNN branch structurally cannot: that gradient comes from an atomic epilogue and is not reproducible against itself, while the tilelang kernel is bitwise stable -- verified bit-identical fused vs eager for out/dq/dkv/d_sink.
PR Category
Distributed Strategy
PR Types
New features
Description
增加mqa tilelang的backend
是否引起精度变化
否