[release/0.4] Fuse HCA inverse RoPE into the VHA postmix, and the sparse-attn KV index remap

[DanielSun11]

PR Category

Performance Optimization

PR Types

New features, Performance

Description

Merged in dev: #1772

将已合入 develop 的 1772（commit 99e711822a70f8e338a4bfae7b28f750a541fff9）cherry-pick 到 release/0.4。两个融合开关默认均为 False，未开启时 release/0.4 的行为与本 PR 之前完全一致。

一、带来的改动

1. fuse_inv_rope_into_vha_postmix：把 HCA 层的 inverse RoPE 折进 ungrouped VHA postmix 的 [nh,nh] GEMM（triton_ops/inv_rope_vha_postmix_fusion.py），不再物化全宽的 inv_rope(O)。DSv4HybridAttention._can_fuse_inv_rope_postmix 对所有不适用的组合（开关关闭、vha_postmix_grouped=True、apply_rope_fusion=False、high_precision_rope=True、postmix 自带 selective recompute 包装）一律回退而不是报错。
2. sparse_attn_global_kv_idx_remap_fusion：把「per-batch-local → flat-global」KV 列索引重映射（idx + b * seqlen_kv）的七个 elementwise kernel 合成一个 Triton kernel（triton_ops/local_to_global_idxs_fusion.py），结果逐位相同。覆盖 CompressedSparseAttention 与 MQALatentAttention 在 "cudnn" 后端下的前向与反向；"tilelang" / "unfused" 后端不受影响。
3. _derive_csa_doc_boundaries 去 host-sync：is_boundary[0] = True 的 1 字节 pageable cudaMemcpy 换成 concat（原本会阻塞 host 直到 device queue 排空，在 layer43 配置上落在 DeepEP combine 之后，每个 -2 层约 2.9 ms）；单行长序列上的 paddle.cummax 换成 cumsum + gather（seqlen 131072 时 5.3 ms → 约 0.04 ms）。

二、冲突解决

release/0.4 上有 develop 尚未合入的 padded-head / sub-512-latent 稀疏 attention 路径，因此产生两处冲突。两处都是保留 release 独有逻辑、只把 dev 的改动叠加上去：

1. src/paddlefleet/fusions/csa_sparse_attn.py
   • 保留 release 的 kernel 侧 head 数 kh 用于 lse_flat（dev 那边仍是 np_heads），保留 release 由 ctx.compacted_idxs 驱动的 topk_length 重算；
   • 采纳 dev 的 local_to_global_flat(topk_idxs, s_kv, fused=ctx.global_kv_idx_remap_fusion) 分发；
   • docstring 是纯 additive 冲突，同时保留 dev 新增的参数说明和 release 的 head-tile / latent-padding 说明。
2. tests/single_card_tests/ai_edited_test/fusions/test_csa_sparse_attn_backends.py
   • 保留 release 伪 ctx 上的 compacted_idxs 字段。dev 在这个文件里新增的 global_kv_idx_remap_fusion=False 在 release 上已经存在，所以该文件最终没有改动（本 PR 是 19 个文件，dev 侧是 20 个）。

三、验证

• 逐 hunk 比对了全部 20 个文件的 patch body，除上述两个冲突文件外，其余 18 个文件与 dev 侧改动完全一致。
• 核对了 CSASparseAttention.forward 的形参顺序与 apply 的位置传参一致，global_kv_idx_remap_fusion 在末位。
• 排查了所有用 SimpleNamespace 伪造 ctx 直接调 backward 的测试，没有遗漏 global_kv_idx_remap_fusion 属性的（否则新增的 ctx.global_kv_idx_remap_fusion 读取会抛 AttributeError）。release 上另外两个伪 ctx 测试本来就带了这个字段。
• 全部改动文件 py_compile 通过；pre-commit run --files <changed> 全绿（ruff check / ruff format / typos / copyright / whitespace）。
• 新增的三个 Triton kernel 测试需要 GPU，本地未跑，交给 CI。

首轮 CI 里 Unit test (multi-card) 与 Integration test (H20, single card) 的失败发生在 Clone / Install 阶段，日志为 You are using GPU version Paddle, but your CUDA device is not set properly 且 paddlefleet_ops/__init__.py:114 的 paddle.cuda.get_device_capability() 抛 Place(cpu)，即 runner 上没有可见 GPU，与本次改动无关，需重跑。

是否引起精度变化

否

两个开关默认关闭；开启后前向、激活梯度以及 postmix U/V 的梯度都与不融合版本逐位一致，由 dev 侧新增的单测钉住。第 3 项 _derive_csa_doc_boundaries 的改写是等价重写（concat 与 cumsum + gather 产出的 boundary 与 doc_start_per_pos 与原实现相同），不涉及数值。

[codecov-commenter]

Codecov Report

✅ All modified and coverable lines are covered by tests.
⚠️ Please upload report for BASE (release/0.4@49c3dc9). Learn more about missing BASE report.

Additional details and impacted files

@@               Coverage Diff               @@
##             release/0.4     #1843   +/-   ##
===============================================
  Coverage               ?   100.00%           
===============================================
  Files                  ?         8           
  Lines                  ?        42           
  Branches               ?         9           
===============================================
  Hits                   ?        42           
  Misses                 ?         0           
  Partials               ?         0           

Flag	Coverage Δ
coverage_combine	100.00% <100.00%> (?)

Flags with carried forward coverage won't be shown. Click here to find out more.

Files with missing lines	Coverage Δ
...efleet/cudnn_ops/attn/csa_sparse_attn_fwd_cudnn.py	100.00% <100.00%> (ø)
src/paddlefleet/fusions/csa_sparse_attn.py	100.00% <100.00%> (ø)
src/paddlefleet/fusions/csa_sparse_attn_utils.py	100.00% <100.00%> (ø)
src/paddlefleet/fusions/mqa_sparse_attn.py	100.00% <100.00%> (ø)
src/paddlefleet/transformer/csa_attention.py	100.00% <100.00%> (ø)
...c/paddlefleet/transformer/dsv4_hybrid_attention.py	100.00% <100.00%> (ø)
...rc/paddlefleet/transformer/mqa_latent_attention.py	100.00% <100.00%> (ø)
src/paddlefleet/transformer/transformer_config.py	100.00% <100.00%> (ø)

🚀 New features to boost your workflow:

• ❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.