support mqa indexer backward determine - #1683
Conversation
a72eff8 to
c0572fb
Compare
There was a problem hiding this comment.
| 序号 | 位置 | 优先级 | 状态 |
|---|---|---|---|
| 1 | transformer_config.py |  |
✅ |
| 2 | mqa_latent_sparse_bwd.py | |
✅ |
| 3 | mqa_latent_attention.py | |
✅ |
| 4 | mqa_latent_sparse_bwd.py | |
✅ |
| 5 | mqa_latent_sparse_bwd.py | |
✅ |
当前 head 7339b4a 已对小于 kernel tile 的 head 维度进行一致 padding，并在返回前裁剪；H=8 和 H=24 的新增测试覆盖了实际 kernel 与参考实现的梯度对齐。
| """ | ||
|
|
||
| stage1_overlap: bool = False | ||
| mqa_sparse_attn_backward_backend: str = "cudnn" |
There was a problem hiding this comment.
这里把已有的 stage1_overlap: bool = False 直接替换掉了。core_transformer_config_from_args 只遍历 dataclasses.fields，所以现有 YAML/CLI 的 stage1_overlap 会被静默丢弃，而直接 TransformerConfig(stage1_overlap=...) 也会变成 TypeError；transformer_layer.py:87 仍按该字段决定共享 MTP 的 sharding overlap。请保留原字段，并在其后新增 mqa_sparse_attn_backward_backend。
| be reused across all chunks (``bwd_det`` overwrites every slot, so the buffer | ||
| needs no clearing between launches). | ||
| """ | ||
| row_bytes = topk * dk * 4 |
There was a problem hiding this comment.
row_bytes 少计算了 batch 维度：实际在下方分配的 dkv_buf 形状是 [b, sc, topk, dk]，预算应按 b * sc * topk * dk * 4 计算。以默认预算和文档中的 S=8192, L=640, Dk=576 为例，B=1 会选择 sc=8192 并占用约 11.25 GiB；同样合法的 B=2 仍选择这个 chunk，却分配约 22.5 GiB，直接违背 12 GiB 上限并可能 OOM。请将 b 纳入 _pick_chunk 的计算（或在分配前按实际 shape 重新限制 chunk）。
| topk_probs, | ||
| loss_coeff, | ||
| "cudnn", | ||
| self.indexer_backend, |
There was a problem hiding this comment.
这里把 self.indexer_backend 传给了 MQA sparse loss scaler，但配置校验允许 csa_indexer_backend="unfused"。该 scaler 的 backward 只处理 "cudnn" 和 "tilelang"，其它值会在 csa_attention.py 的 NotImplementedError 分支失败；因此 hybrid_mla_attention="mqa_dsa" 的合法 unfused 配置会在第一次 sparse backward 时崩溃。请为该路径实现 unfused backward，或在 MQA sparse 模式下明确拒绝/改用受支持的后端。
| _, s_kv, kv_dk = kv.shape | ||
| dv = grad_out.shape[-1] | ||
|
|
||
| assert kv_dk == dk, f"kv width {kv_dk} != query width {dk}" |
There was a problem hiding this comment.
这组输入形状和 kernel 约束使用了 Python assert（同一函数的 196-204 行）。在 python -O 下这些检查会被完全移除，非法的 Dk/Dv、head 数或 block_h 会继续进入 TileLang JIT 和显存分配，导致更晚的 opaque 错误甚至越界风险。请改为显式 ValueError/参数校验；前向新增的同类 assert 也应一并处理。
c0572fb to
70fe9ef
Compare
risemeup1111
left a comment
risemeup1111
left a comment
There was a problem hiding this comment.
当前 head 仍有以下未解决的 P1，请按首评行间评论修复后再请求复核：
- mqa_latent_sparse_bwd.py：
dkv_buf显存预算仍未计入 batch 维度。 - mqa_latent_attention.py：合法的
unfusedindexer backend 仍会在 sparse backward scaler 中失败。 - mqa_latent_sparse_bwd.py：运行时形状校验仍使用 Python
assert。
首评 Board 已更新：stage1_overlap 已解决，其余三项仍需处理。
93ee790 to
f20f558
Compare
f20f558 to
1e1ad70
Compare
risemeup1111
left a comment
risemeup1111
left a comment
There was a problem hiding this comment.
当前 head 新增的 TileLang backward 在每卡 H<16 时仍有未解决的 P1：bwd_det 会把 head 数向上 pad 到 16，但调用侧仍传入未 padding 的 8-head 张量。详情见对应行间评论。
``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.
Codecov Report❌ Patch coverage is Additional details and impacted files@@ Coverage Diff @@
## release/0.4 #1683 +/- ##
==============================================
Coverage ? 94.31%
==============================================
Files ? 4
Lines ? 176
Branches ? 34
==============================================
Hits ? 166
Misses ? 7
Partials ? 3
Flags with carried forward coverage won't be shown. Click here to find out more.
🚀 New features to boost your workflow:
|
4 participants
PR Category
Execute Infrastructure
PR Types
New features
Description
增加mqa tilelang的backend
是否引起精度变化
否
Merged in dev: #1776