[New features] Fuse HCA inverse RoPE into the VHA postmix, and the sparse-attention KV index remap - #1772
Open
DanielSun11 wants to merge 8 commits into
Conversation
…arse-attention KV index remap Two independent, bitwise-identical Triton fusions on the DSv4-hybrid path, both behind config switches that default to off. 1. fuse_inv_rope_into_vha_postmix The HCA inverse RoPE materialises a full-width inv_rope(O) only to feed it to the ungrouped VHA postmix [nh,nh] GEMM, costing an extra read+write of the whole attention output and a second live copy of it. RoPE only touches the trailing qk_pos_emb_head_dim channels while the GEMM contracts the head axis, so the channel axis is a pure N dimension and the result can be assembled from a full-width GEMM on the unrotated output plus a narrow GEMM on the rotated pe channels -- no wide intermediate. The weight gradient contracts over the full channel axis, so it cannot be split; it needs head-major operands, which is what Paddle's own matmul backward spends 73% of its time building with two full-size TilingSwapDim1And2 passes (2.5 TB/s). Producing them directly lets the RoPE ride along in the same pass and runs both at the HBM roof: 324 us and 320 us at 6.6/6.7 TB/s against 1965 us and 1629 us. Measured on one B30Z (sm10.3), 64K context with context_parallel_size=4, i.e. sq=16384 and a 1 GiB attention output, nh=64 / nope=448 / pe=64: forward 643.7 -> 459.4 us (-28.6%) backward 2699.4 -> 1828.9 us (-32.3%) fwd+bwd peak memory 5.00N -> 4.00N (-1024 MiB), live after forward -1028 MiB 37 HCA layers with full recompute: -45.8 ms/step Non-power-of-two head counts fall back to matmul_grad: the degenerate [nh,nh] x K GEMM makes cuBLAS pick its algorithm unpredictably and the head-major operands stop matching bit for bit (seen at nh=3). 2. sparse_attn_global_kv_idx_remap_fusion The per-batch-local -> flat-global KV column index remap consumed by the cuDNN / FlashMLA sparse-attention kernels spends seven elementwise kernels on the full [b*sq, topk] table; the Triton kernel does it in one, bit-identically. Applies to both CompressedSparseAttention and MQALatentAttention on the "cudnn" backend. Testing New: tests/single_card_tests/test_inv_rope_vha_postmix_fusion.py (32), tests/single_card_tests/transformer/test_inv_rope_vha_postmix_layer.py (8), tests/single_card_tests/test_local_to_global_idxs_fusion.py (24). The layer-level test builds a real DSv4HybridSelfAttention, runs it with the switch off and on and requires the output, the input gradient and every parameter gradient to be bitwise equal; it also asserts the gate really flips, moves the zero-initialised postmix V off identity (otherwise every postmix gradient is exactly zero and the weight-gradient comparison is vacuous) and checks that two unfused runs agree before comparing anything. The cuBLAS and Paddle properties the fusion leans on -- split-N GEMM bitwise equality, matmul_grad vs the explicit formulas, the head-major weight gradient across shapes -- are asserted directly, so a library upgrade fails CI instead of silently drifting the loss.
CI on sm90 caught `head-major wgrad (128,4,64)` differing from Paddle's own
matmul backward in 8/16 elements, while the same shape was bitwise equal on
sm10.3 where the fusion was developed. Hand-rolling that GEMM over head-major
operands means cuBLAS selects the algorithm for a small [nh,nh] x K problem, and
that selection is per-architecture -- so the `nh & (nh - 1)` gate cannot make the
route safe and no shape-based gate can.
The backward now always rebuilds the rotated tensor with rope_full_out_of_place
and hands both gradients to paddle._C_ops.matmul_grad, i.e. the very op the
unfused postmix GEMM's backward calls. That is bitwise correct by construction
everywhere. Drops _rope_transposed_kernel / rope_full_to_transposed and
_transpose_heads_first_kernel / transpose_heads_first, which existed only to feed
the removed route.
Cost at the production shape (sq=16384, 1 GiB attention output), one B30Z:
forward 642.8 -> 459.4 us (-28.5%, unchanged)
backward 2700.1 -> 3012.3 us (+312 us: the full-width RoPE moves into the
backward instead of being folded into a
transpose that was already required)
37 HCA layers with full recompute: -2.0 ms/step (was -45.8 ms/step)
live after forward still -1028 MiB/layer, but the fwd+bwd peak is now neutral
(5.00N both) because the rebuilt tensor is live alongside dOut and dO
The forward split is untouched and still the only cuBLAS property relied on;
test_split_n_gemm_is_bitwise now sweeps ten shapes including the small head
counts that broke the other route, and the end-to-end bitwise comparison gained
TestSmallHeadCounts covering (128,4,64) and friends.
risemeup1111
approved these changes
Aug 18, 2026
Both features added something the hand-built stubs in ai_edited_test/ do not
know about, so 21 tests across four files failed on the new interface rather
than on any behaviour:
- `_full_attn_forward` now asks `self._can_fuse_inv_rope_postmix(...)` before
the inverse-RoPE block, and the DSv4 recompute tests drive that method on a
SimpleNamespace. Bind the real gate (plus the `vha_postmix_grouped` flag it
reads) instead of a lambda, so a change to its conditions still surfaces
there; with `use_vha_postmix` False it returns False and those tests stay on
the unfused path they were written for.
- `global_kv_idx_remap_fusion` is now stashed on the sparse-attention ctx and
passed to the cuDNN forward, so `_FakeCtx`, the SimpleNamespace ctx and the
`fake_fwd` signature need it.
All four files pass locally: 23, 21 (+6 subtests), 19, 18 (+1 skipped).
risemeup1111
approved these changes
Aug 18, 2026
The single-card CI stage has hit its 120 minute budget three runs in a row, each time stalling inside test_local_to_global_idxs_fusion.py with no progress output for ~58 minutes. The file takes 24s locally, so nothing here reproduces off the CI runner -- but it was also by far the most host-memory hungry file in the suite, and that is worth removing whether or not it is the cause. test_ernielite_real_shapes built two [1, 16384, 2176] tables (35.7M entries) per dtype, each with a float64 pad mask of the same shape, then pulled both the eager and the fused result back to host for a numpy comparison: on the order of 2-3 GB of host allocation for one test. test_ernielite_hca_shape did the same at 10.5M entries and is the test the stall sits on. The kernel is one program per row with cdiv(topk, BLOCK_K) along topk, so the row count is only a grid multiplier and exercises no additional path. Trim sq from 16384 to 1024 in those two tests, keeping the topk widths (640 / 2176) and the seqlen_kv that sets the int32 offset range -- the properties they are actually about. Draw the pad mask as int16 rather than float64 while here. Also add a one-row one-column launch as the first test in the file, so the next CI run distinguishes "the first Triton compile stalls" (no progress at all) from "a later test's volume is the problem" (that dot appears, then the stall). File drops from 24.3s to 15.8s locally and from ~3 GB to ~150 MB of peak host allocation; 25 tests pass.
risemeup1111
approved these changes
Aug 18, 2026
The three TestConfigSwitch end-to-end tests reach FlashMLA / DSA kernels that are SM100-only -- the backward their docstrings quote lives in sparse_attention_backward/dsa_bwd_sm100.py -- but they guard on `import paddlefleet_ops.flash_mla` succeeding. That is a "is the library present" check, not a capability check, and on the CI runner the library *is* present: it is SM 9.0 (runner group fleet-h-single-card) and the log says "Successfully loaded ecosystem library: flash_mla". So the guard never fires and the tests call SM100 kernels on Hopper, where the behaviour is undefined. test_fused_sink_grad.py already knows this -- it drives the backward through a _FakeCtx specifically because "the CI runner is Hopper; the kernels are SM100+" -- so gate the same way test_hysparse_online_tilelang_train_step gates its TileLang kernels: check get_device_capability() and skip below SM 10.x, keeping the import check as a second step. This is correct on its own merits (calling SM100 kernels on SM 9.x tests nothing) and it removes the most hardware-sensitive thing in a file that has now stalled the single-card stage three runs in a row. It is not proven to be the stall: the stall appears to sit on the first test in the file, which touches none of this. The one-row launch added as the first test in the previous commit is what will settle that on the next run. 25 tests pass locally (SM 10.3); with the capability faked to (9, 0) the guard skips instead of entering the kernels.
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.
Contributor
There was a problem hiding this comment.
| 序号 | 位置 | 优先级 | 规则来源 | 状态 |
|---|---|---|---|---|
| 1 | SM 兼容性门禁 |  |
仓库规则：硬件保护与兼容性 | 🟡 |
Powered by Nyanpasu with gpt-5.6-sol 默认推理级别, please check the suggestions carefully.
develop's PaddlePaddle#1762 (dual-chunk indexer CP rebalance) appended both a new ``indexer_dualchunk`` attribute and three helper methods at the exact end of ``MQALatentAttention.__init__`` where this branch had added ``global_kv_idx_remap_fusion``. Both additions are independent, so both are kept.
risemeup1111
left a comment
risemeup1111
left a comment
Contributor
There was a problem hiding this comment.
本轮同步提交为合并基线后的增量，未新增可确认问题；现有 P3 仍由 canonical thread 跟踪，状态保持未解决。
Powered by Nyanpasu with gpt-5.5 xhigh, please check the suggestions carefully.
Codecov Report✅ All modified and coverable lines are covered by tests. Additional details and impacted files@@ Coverage Diff @@
## develop #1772 +/- ##
===========================================
Coverage ? 100.00%
===========================================
Files ? 8
Lines ? 42
Branches ? 9
===========================================
Hits ? 42
Misses ? 0
Partials ? 0
Flags with carried forward coverage won't be shown. Click here to find out more.
🚀 New features to boost your workflow:
|
LiYuRio
approved these changes
Aug 19, 2026
Contributor
|
@risemeup1111 重新review |
Contributor
|
已按当前 head |
Contributor
Author
这个单测不能按照你说的改！理论上单测需要保证向后兼容，更高的硬件也需要支持这个单测 |
Contributor
|
@risemeup1111 这个单测不能按照你说的改！理论上单测需要保证向后兼容，更高的硬件也需要支持这个单测 |
Contributor
|
你说得对。我重新核对后确认这里的契约是支持 SM 10 及以上，Triton remap 会在更高架构上 JIT 编译；我此前把当前扩展的预编译架构列表误当成了测试的永久上限。原 P3 不成立，我会在 canonical thread 撤回并更新 Board。 |
risemeup1111
approved these changes
Aug 19, 2026
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
Performance Optimization
PR Types
New features, Performance
Description
在 DSv4-hybrid attention 路径上引入两个相互独立的 Triton 融合，均由默认关闭的 config 开关控制，且与被替换的原路径逐位等价。
一、
fuse_inv_rope_into_vha_postmixHCA 层的 inverse RoPE 会物化一份全宽的
inv_rope(O)，仅仅为了交给 ungrouped VHA postmix 的[nh,nh]GEMM。代价是把整个 attention 输出多读写一遍，并且让它的第二份全宽副本一直存活。RoPE 只作用于每个 head 尾部的
qk_pos_emb_head_dim个通道，而 postmix GEMM 收缩的是 head 轴，所以 channel 轴对这个 GEMM 而言只是纯粹的 N 维：沿 N 轴拆分不改变累加顺序，因此结果可以由「对未旋转输出做一次全宽 GEMM」加上「对已旋转的 pe 通道做一次窄 GEMM」拼出来，全宽中间张量从此不需要存在
dM[h,h'] = sum_{t,c} dOut[t,h,c] * O_roped[t,h',c]的收缩轴含完整 channel 轴，不能拆成部分和，因此必须有一份全宽的旋转张量交给
paddle._C_ops.matmul_grad——也就是不融合路径的 postmix GEMM 反向所调用的同一个 op，这样两个梯度由构造保证在任何架构上都逐位一致。
早先的版本自己用 head-major
[H, B*S, D]操作数手写这个 GEMM（matmul_grad内部正是用两趟2.5 TB/s 的
TilingSwapDim1And2构造这两个 buffer），把 RoPE 折进同一趟，快约 870 us/层，在 sm10.3 上扫过的每个 shape 都逐位相等——但 sm90 上不成立：CI 抓到
head-major wgrad (128,4,64)有 8/16 个元素不同。小的[nh,nh] x KGEMM 上 cuBLAS 按架构选算法，任何基于 shape 的门控都无法让这条路安全，因此已删除。
前向的 channel 拆分是唯一还依赖 cuBLAS 性质的地方（
matmul(M, X[..., a:b]) == matmul(M, X)[..., a:b]），单测里按十种 shape 扫过、包含把另一条路打挂的那些小 head 数。所有不适用的组合——开关关闭、use_vha_attention=False、vha_postmix_grouped=True、apply_rope_fusion=False、high_precision_rope=True、postmix 自带 selective recompute 包装、U/V被冻结——都是回退而不是报错。二、
sparse_attn_global_kv_idx_remap_fusioncuDNN / FlashMLA 稀疏 attention kernel 所消费的「per-batch-local → flat-global」KV 列索引重映射（
idx + b * seqlen_kv），eager 版本要在完整的[b*sq, topk]表上跑七个 elementwise kernel（full+greater_equal+arange+expand+scale+add+where），Triton kernel 一趟做完，结果逐位相同。覆盖所有_local_to_global_flat调用点——CompressedSparseAttention（HCAratio=128与 CSA/DSA1 < ratio < 128）和MQALatentAttention在"cudnn"后端下的前向与反向。对"tilelang"/"unfused"后端无影响。三、测试
新增：
tests/single_card_tests/test_inv_rope_vha_postmix_fusion.py（29 个）tests/single_card_tests/transformer/test_inv_rope_vha_postmix_layer.py（8 个）tests/single_card_tests/test_local_to_global_idxs_fusion.py（24 个）层级测试会构造一个真实的
DSv4HybridSelfAttention，用同一 seed 建两份、只翻开关，要求层输出、输入梯度、以及全部参数梯度（14 个参数中 10 个带梯度）逐位相等。它同时断言门控确实在两次运行之间发生了切换；把零初始化的 postmixV挪离恒等（否则所有 postmix 梯度精确为 0，权重梯度的比对形同虚设）；并要求两次不融合的运行先自行一致，才开始做任何比对。融合所依赖的库行为都被直接断言而非假定，这样上游库升级会让 CI 失败，而不是让 loss 曲线悄悄漂移：沿 N 轴拆分 GEMM 的逐位等价（十种 shape，含小 head 数）、
matmul_grad与显式 dgrad 对 autograd 的一致，以及TestSmallHeadCounts在(128,4,64)等小 shape 上的端到端逐位比对——正是这一类断言在 sm90 上抓出了手写 head-major wgrad 的问题。本分支上重跑的既有测试全部通过：
test_vha_dsv4（38）、test_dsv4_hybrid_attention（96）、test_hca_csa_independent_rope（17）、test_train_indexer_only（24）、test_csa_loss_mask（11）、test_mla_rope_inplace_fusion（8）。是否引起精度变化
否
两个开关默认均为
False；开启后每条路径都与不融合版本逐位一致——前向、激活梯度、以及 postmixU/V的梯度——已在生产 shape 上用零容差比对验证，并由上述单测钉住。