[BugFix] Pad the tilelang MQA backward heads up to the kernel tile - #1784
Merged
Merged
Conversation
``bwd_det`` derives its own tile from the ``H`` it is handed -- ``padded_H = max(next_power_of_2(H), 16)``, ``block_H = min(64, padded_H)`` -- and then copies ``Q[b, s, 0:block_H, :]`` and stores ``dQ[b, s, 0:block_H, :]``. Handing it a head group narrower than 16 does not make it move fewer heads: at ``block_h = 8`` it still moves 16, i.e. it reads and writes 8 heads past the end of an 8-head tensor. Restrict the group width to the two values for which the tile matches (16 / 32) and zero-pad the head axis up to a multiple of it instead. Pad heads are inert: q and dO are zero there, so dP and both dKV terms vanish and Delta is zero, so d_sink is too; their dq rows are sliced off before returning. ``h = 8`` -- the per-rank count of the hybrid-MLA fixtures and of any TP > 4 split -- therefore keeps working without touching the kernel.
Contributor
There was a problem hiding this comment.
Pull request overview
该 PR 修复了 TileLang 确定性 MQA latent sparse backward（mqa_latent_sparse_bwd）在每卡 head 数小于 kernel tile（例如 H=8）时会触发越界读写、进而导致显存破坏的问题（PR 描述中指出源于 #1776 的行为）。
Changes:
- 将
block_h约束为 kernel 自身可安全支持的两种 tile 宽度（16/32），并对不满足分组整除的 head 维度在 wrapper 侧进行零填充与返回前切片。 - 新增/调整单测：覆盖非法
block_h拒绝、H<tile场景的 padding 行为与返回 shape 校验，并收紧_pick_block_h的返回范围验证。
Reviewed changes
Copilot reviewed 2 out of 2 changed files in this pull request and generated no comments.
| File | Description |
|---|---|
| tests/single_card_tests/transformer/test_mqa_latent_sparse_bwd.py | 增补用例，确保非法 tile 宽度被拒绝，并验证 H=8/24 走 padding 路径且返回 shape/数值与参考一致 |
| src/paddlefleet/tilelang_ops/attn/mqa_latent_sparse_bwd.py | 将 block_h 限定为 16/32，新增 head 维 padding + 返回前切片，避免小 head 数触发 kernel 越界读写 |
💡 Add a code-review agent skill or configure MCP servers for context-aware, tailored reviews. Learn more in the docs.
risemeup1111
approved these changes
Aug 19, 2026
Codecov Report❌ Patch coverage is
Additional details and impacted files@@ Coverage Diff @@
## develop #1784 +/- ##
==========================================
Coverage ? 92.00%
==========================================
Files ? 1
Lines ? 25
Branches ? 4
==========================================
Hits ? 23
Misses ? 1
Partials ? 1
Flags with carried forward coverage won't be shown. Click here to find out more.
🚀 New features to boost your workflow:
|
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.
PR Category
Distributed Strategy
PR Types
Bug fixes
Description
修复 #1776 引入的
mqa_latent_sparse_bwd在每卡 head 数小于 16 时的越界访问（release/0.4 侧的同一修复见 #1683）。问题：
bwd_det的 tile 不由调用方决定，它自己算padded_H = max(next_power_of_2(H), 16)、block_H = min(64, padded_H)（sparse_mqa_bwd.py:352-353），然后无条件T.copy(Q[b, s, 0:block_H, :])并T.copy(dQ_shared, dQ[b, s, 0:block_H, :])。原来的_pick_block_h会为h = 8选出 group 宽度 8，kernel 却仍按 16 个 head 搬运：在 8-head 的Q/Lse/Delta上越界读，并且越界写dQ—— 而dq只按[b, s, 8, dk]分配，属于显存破坏。h = 8是 hybrid-MLA 单测 fixture 以及 TP > 4 每卡的 head 数，h = 24 / 40同样会落到小于 16 的宽度。修法（不改 kernel）：把 group 宽度限制为
bwd_det的 tile 恰好等于入参的两个值 16 / 32，H不是其倍数时由 wrapper 在 head 轴零填充，跑完再切回。_pick_block_h(h)→32 if h % 32 == 0 else 16；显式传入 8 / 48 / 64 直接ValueError；h_pad = ceil(h / block_h) * block_h，对query/grad_out/delta/lse_tl补 0、attn_sink补-1e30；dq按h_pad分配，返回前切成[b, s, h, dk]，d_sink切成[h]。填充 head 是惰性的：q 与 dO 为 0 ⇒
dP = P * (dO·KV - Delta) * scale = 0，dKV的dP^T·Q与P^T·dO两项皆为 0；Delta = 0⇒d_sink = -Delta * exp2(...) = 0。只有填充 head 的dq行无意义，返回前被切掉。测试：新增
test_only_the_kernels_own_tile_widths_are_accepted（8 / 48 / 64 必须报错）与test_head_count_below_the_tile_is_padded_not_truncated（h = 8与h = 24走填充路径，校验返回 shape 精确为[s, h, DK]/[h]且与 fp32 参考一致）；_pick_block_h的用例改为遍历1..64断言结果只能是 16 或 32。本地实测：test_mqa_latent_sparse_bwd.py25 passed / 119 subtests、test_mqa_latent_attention.py64 passed、test_fused_sink_grad.py21 passed。注：原先
h = 8的用例能通过是因为它只比较前h个 head 的梯度，而越界写落在分配之外的显存里，既不改变被比较的数值也不必然崩溃 —— 数值断言对这类问题无效，因此改为同时校验 shape 与拒绝非法 tile 宽度。是否引起精度变化
否
backward_backend默认仍为cudnn，不走本文件；tilelang 路径在h为 16 / 32 的倍数时行为与修复前一致（同样的 group 宽度、同样的求和顺序），仅h不是 16 倍数时从「越界」变为「填充后正确」。