[New features] Add HySparse MQA/MLA sparse attention integration#1453
[New features] Add HySparse MQA/MLA sparse attention integration#1453ForFishes wants to merge 11 commits into
Conversation
|
|
There was a problem hiding this comment.
Pull request overview
该 PR 将 HySparse（块级打分 + TopK 选块 + 块稀疏 gather，arXiv:2602.03560）的算子与模型侧 MLA-absorbed MQA 注意力路径打通，并通过 enable_hy_sparse_attention 配置开关启用，用于在 full-attention 层生成共享 KV + TopK block 索引，并在 SWA 层复用以执行块稀疏注意力。
Changes:
- 新增/扩展 TileLang HySparse 相关算子（MQA/MHA block-score、MQA block-sparse、SWA wrapper）以及对应的 PyLayer 可微封装与参考实现。
- 在
MQASelfAttention/HySparseTransformerLayer/gpt_layer_specs中完成模型侧接线：full 层产出共享 KV + block 索引，SWA 层消费并叠加 sparse 分支输出。 - 增加多组算子精度/反传测试与网络集成测试（含 stop_gradient 合约、非连续梯度等）。
另外：PR 标题当前不符合仓库要求的 "[CLASS]Title" 格式；建议改为类似 [Feature] Add HySparse MQA/MLA sparse attention integration（或按项目约定选择合适的 CLASS）。
Reviewed changes
Copilot reviewed 21 out of 21 changed files in this pull request and generated 3 comments.
Show a summary per file
| File | Description |
|---|---|
| tests/single_card_tests/test_mqa_self_attention.py | 新增网络级集成测试，覆盖 MLA vs MQA 对齐与跨层 KV sharing 行为 |
| tests/single_card_tests/custom_ops/test_hysparse_swa.py | 新增 SWA MQA wrapper 的前后向精度与 autograd 合约测试 |
| tests/single_card_tests/custom_ops/test_hysparse_mha_block_score.py | 新增 MHA（per-head KV）block-score 的参考对齐与 bwd 覆盖 |
| tests/single_card_tests/custom_ops/test_hysparse_differentiable_ops.py | 新增 PyLayer 包装的 stop_gradient / 非连续梯度 / 正确性测试 |
| tests/single_card_tests/custom_ops/test_hysparse_block_attn.py | 扩展 MQA 测试以覆盖 Dk!=Dv 的 absorbed-MLA 形状 |
| src/paddlefleet/transformer/transformer_layer.py | 引入 HySparseTransformerLayer，补齐 shared_kv 透传与 MTP split/restore |
| src/paddlefleet/transformer/transformer_config.py | 增加 HySparse 配置项：enable 开关、block_size、topk |
| src/paddlefleet/transformer/multi_latent_attention.py | full-attention 路径接入 FA4 fused block-score + TopK，并产出 shared_kv |
| src/paddlefleet/tilelang_ops/hysparse/swa_attn.py | 新增 SWA（滑窗）MQA 注意力 wrapper（复用 block-score kernel） |
| src/paddlefleet/tilelang_ops/hysparse/reference.py | 扩展参考实现：支持 Dv 维度、补充 MHA block-score reference、滑窗 valid_range |
| src/paddlefleet/tilelang_ops/hysparse/pipeline.py | TopK block 选择前对输入 detach，避免 TopK 进入反传导致崩溃 |
| src/paddlefleet/tilelang_ops/hysparse/differentiable_ops.py | 新增 block-score / block-sparse 的 PyLayer 可微封装与 pipeline 便捷函数 |
| src/paddlefleet/tilelang_ops/hysparse/block_sparse_attn_mqa.py | 扩展 block-sparse MQA fwd 以支持 Dk!=Dv，并加强 shape 校验 |
| src/paddlefleet/tilelang_ops/hysparse/block_sparse_attn_mqa_bwd.py | 扩展 block-sparse MQA bwd 以支持 Dk!=Dv，并调整 shared-mem 预算逻辑 |
| src/paddlefleet/tilelang_ops/hysparse/block_score_fa4.py | 新增 FA4 fused block-score wrapper（含环境探测/快速失败） |
| src/paddlefleet/tilelang_ops/hysparse/block_score_attn.py | 扩展 MQA block-score fwd 以支持 Dk!=Dv，并加强 shape 校验 |
| src/paddlefleet/tilelang_ops/hysparse/block_score_attn_mha.py | 新增 MHA（per-head KV）block-score fwd kernel 与 host 侧接口 |
| src/paddlefleet/tilelang_ops/hysparse/block_score_attn_mha_bwd.py | 新增 MHA block-score bwd kernel 与 host 侧接口 |
| src/paddlefleet/tilelang_ops/hysparse/block_score_attn_bwd.py | 扩展 MQA block-score bwd 以支持 Dk!=Dv，并调整 shared-mem 预算逻辑 |
| src/paddlefleet/tilelang_ops/hysparse/init.py | 导出新增算子/封装 API（含 SWA / FA4 wrapper） |
| src/paddlefleet/models/gpt/gpt_layer_specs.py | 启用 HySparse 时切换 attention/layer 类型到 MQASelfAttention/HySparseTransformerLayer |
| if self.config.enable_hy_sparse_attention and shared_kv is not None: | ||
| # Compressed KV latent shared with block-sparse attention in SWA | ||
| # layers (single MQA head): [B, S, 1, kv_lora_rank + qk_rope_head_dim]. | ||
| shared_key = paddle.concat( | ||
| [kv_compressed.unsqueeze(2), k_pos_emb], axis=-1 | ||
| ) | ||
| shared_kv.append(shared_key) | ||
| # block_indices produced by the MHA block-score path above. | ||
| shared_kv.append(block_indices) |
| ctx.save_for_backward(q, k, v, startend_row_indices, out, lse) | ||
| ctx.sm_scale = sm_scale | ||
| ctx.causal = causal | ||
| return out, lse | ||
|
|
||
| @staticmethod | ||
| def backward(ctx, dout, *_): | ||
| from paddlefleet_ops.flash_mask.cute.flashmask_utils import ( | ||
| FlashMaskInfoPaddle, | ||
| ) | ||
| from paddlefleet_ops.flash_mask.cute.interface import _flash_attn_bwd | ||
|
|
||
| q, k, v, startend_row_indices, out, lse = ctx.saved_tensor() | ||
| flashmask_info = None | ||
| if startend_row_indices is not None: | ||
| flashmask_info = FlashMaskInfoPaddle( | ||
| startend_row_indices=startend_row_indices, | ||
| is_causal=ctx.causal, | ||
| ) |
| # Grads for the tensor inputs that require grad: q, k, v only. | ||
| return dq, dk, dv |
This comment was marked as outdated.
This comment was marked as outdated.
Sorry, something went wrong.
PaddleFleet Log Analysis
日志分析报告
失败的测试case: 根本原因分析: 修复建议:
🔄 每次 Re-run 后自动更新 |
ForFishes
changed the title
This comment was marked as outdated.
This comment was marked as outdated.
Sorry, something went wrong.
risemeup1111
left a comment
risemeup1111
left a comment
There was a problem hiding this comment.
已复查当前提交，代码层面仍有会导致导入或训练反向执行失败的阻塞问题，具体请看 inline comments。
优先级：P3 非行级：PR 描述模板字段不属于 diff 行，无法挂 inline。当前
Check PR Description仍在报PR Category/PR Types不符合允许枚举，建议按仓库模板在描述开头补充，例如### PR Category使用Performance Optimization，### PR Types使用New features，再保留现有 What/Tests 内容。处理要求：请针对该评论进行回复（同意并已修改请回复 Done，不同意请说明理由）。
| hysparse_forward_mqa, | ||
| select_topk_blocks, | ||
| ) | ||
| from .block_score_fa4 import block_score_fa4_attn_fwd |
There was a problem hiding this comment.
这里在包初始化阶段直接导入 block_score_fa4，会把 paddlefleet_ops.flash_mask.cute.* 的 SM10/Blackwell 依赖提前带入。paddlefleet_ops.__init__ 里只有 is_flash_mask_available() 为 true 时才加载 flash_mask，否则会通过 HardwareIncompatibleBlocker 阻断 paddlefleet_ops.flash_mask 导入；因此在 H20/A100 或非 Blackwell 环境中，只要代码执行 from paddlefleet.tilelang_ops.hysparse import block_sparse_mqa_attention / sliding_window_mqa_attention，即使不走 FA4 full layer，也会在导入阶段失败。
处理要求：请针对该评论修复并提交新的 commit。
建议让 hysparse 包级入口保持非 FA4 路径可导入，把 FA4-only 的 cute import 延迟到实际调用并加可用性检查，例如：
# hysparse/__init__.py: 不要在包初始化时 eager import block_score_fa4
from .differentiable_ops import block_sparse_mqa_attention
from .pipeline import select_topk_blocks
from .swa_attn import sliding_window_mqa_attention
def block_score_fa4_attn_fwd(*args, **kwargs):
from .block_score_fa4 import block_score_fa4_attn_fwd as _impl
return _impl(*args, **kwargs)同时在 block_score_fa4.py 内部用 paddlefleet_ops.is_flash_mask_available() 做明确报错/跳过，或把 FlashMaskInfoPaddle、_flash_attn_fwd/_bwd 的导入移动到 guarded 的执行路径中。
There was a problem hiding this comment.
新提交 99e25e6 又把这个问题带回来了：hysparse/__init__.py 仍执行 from .block_score_fa4 import block_score_fa4_attn_fwd，而当前 block_score_fa4.py 在模块顶层导入 paddlefleet_ops.flash_mask.cute.flashmask_utils/interface。相比上一版，_import_fa4() / is_flash_mask_available() 懒加载也被删除，所以非 SM100 或无 flash_mask 的环境会在导入 paddlefleet.tilelang_ops.hysparse 时先失败，连只使用 sliding_window_mqa_attention 的路径也无法加载。
处理要求：请针对该评论修复并提交新的 commit。
请恢复 package 级入口的懒加载/可用性检查，例如：
# hysparse/__init__.py 不要在模块初始化时导入 FA4-only 依赖
def block_score_fa4_attn_fwd(*args, **kwargs):
from .block_score_fa4 import block_score_fa4_attn_fwd as _impl
return _impl(*args, **kwargs)| ], | ||
| ) | ||
| # Grads for the tensor inputs that require grad: q, k, v only. | ||
| return dq, dk, dv |
There was a problem hiding this comment.
_BlockScoreFA4Attn.forward() 的 Tensor 输入包含 q/k/v/block_logit/startend_row_indices，但 backward 这里只返回了 dq/dk/dv。Paddle 的 PyLayer 会按前向 Tensor 输入位置映射 backward 返回值，少了 block_logit 和 startend_row_indices 的 None 槽位会导致反向执行契约不匹配；另外当前 ctx.save_for_backward(q, k, v, startend_row_indices, out, lse) 在 startend_row_indices=None（该 API 默认允许）时也会把非 Tensor 存入 saved tensors。
处理要求：请针对该评论修复并提交新的 commit。
建议同时修复保存逻辑和返回槽位，形状如下：
ctx.has_mask = startend_row_indices is not None
if ctx.has_mask:
ctx.save_for_backward(q, k, v, startend_row_indices, out, lse)
else:
ctx.save_for_backward(q, k, v, out, lse)
ctx.needs_grad = (
not q.stop_gradient,
not k.stop_gradient,
not v.stop_gradient,
)
ctx.mark_non_differentiable(lse)
...
if ctx.has_mask:
q, k, v, startend_row_indices, out, lse = ctx.saved_tensor()
else:
q, k, v, out, lse = ctx.saved_tensor()
startend_row_indices = None
...
gq, gk, gv = ctx.needs_grad
return (
dq if gq else None,
dk if gk else None,
dv if gv else None,
None, # block_logit
None, # startend_row_indices
)There was a problem hiding this comment.
新提交 99e25e6 回退了这个修复，而且当前 forward 还新增了 block_bos 这个 Tensor 输入：ctx.save_for_backward(q, k, v, startend_row_indices, out, lse) 仍会在 startend_row_indices is None 时保存非 Tensor，backward 仍只 return dq, dk, dv。当前 Tensor 输入至少包含 q/k/v/block_logit/block_bos，有 mask 时还包含 startend_row_indices，因此反向返回槽位还缺 block_logit、block_bos（以及 mask Tensor 对应的 None）。
处理要求：请针对该评论修复并提交新的 commit。
建议恢复“只保存 Tensor + 显式 None 槽位”的形状，并把新增的 block_bos 也纳入返回契约：
ctx.save_for_backward(q, k, v, out, lse)
ctx.startend_row_indices = startend_row_indices
ctx.has_startend = startend_row_indices is not None
ctx.needs_grad = (
not q.stop_gradient,
not k.stop_gradient,
not v.stop_gradient,
)
...
gq, gk, gv = ctx.needs_grad
grads = [
dq if gq else None,
dk if gk else None,
dv if gv else None,
None, # block_logit
None, # block_bos
]
if ctx.has_startend:
grads.append(None) # startend_row_indices
return tuple(grads)| ) | ||
| shared_kv.append(shared_key) | ||
| # block_indices produced by the MHA block-score path above. | ||
| shared_kv.append(block_indices) |
There was a problem hiding this comment.
这里无条件追加 block_indices，但 block_indices 只在上面的 HySparse full-attention 分支里赋值。当前分支顺序是先判断 self.recompute_core_attention and self.training，所以当训练时同时打开 enable_hy_sparse_attention=True 和 recompute_core_attention=True 时，会先走 recompute 的普通 core_attention 分支，然后执行到这里触发 UnboundLocalError。这个配置组合是现有 MLA attention 的可用训练开关，不能因为 HySparse 打开后直接崩溃。
处理要求：请针对该评论修复并提交新的 commit。
建议把 HySparse full-attention 分支放在 recompute 分支之前，或显式禁止这个组合。若选择保留 HySparse，可按下面形状调整：
if self.config.enable_hy_sparse_attention and shared_kv is not None:
core_attn_out, block_indices = self._hy_sparse_full_attention(
query,
key,
value,
attn_mask_startend_row_indices,
)
elif self.recompute_core_attention and self.training:
core_attn_out = recompute(
self.core_attention,
...
)
else:
core_attn_out = self.core_attention(...)如果当前 FA4 full path 暂不支持该 recompute 模式，也请在进入这里前抛出明确的 ValueError，避免后续引用未定义变量。
ForFishes
changed the title
risemeup1111
left a comment
risemeup1111
left a comment
There was a problem hiding this comment.
已复查新提交。上一轮提出的 3 个 P1 代码问题在当前 head 中已看到对应修复；本轮新增发现一个 windowed MQA backward 尾块越界读风险，详情已放在 inline comment。该问题需要修复后再继续复查。
| block_N=block_N, | ||
| block_B=block_B, | ||
| ) | ||
| dq = bwd(q, k, v, do, lse, delta, valid_range, block_range, dk, dv) |
There was a problem hiding this comment.
这里把原始未 padding 的 q / do 直接传入 windowed_mqa_bwd，但 kernel 网格是 T.ceildiv(seq_len, BM)，最后一个 query tile 会执行 T.copy(Q[bb, bm * BM : (bm + 1) * BM, ...]) 和 T.copy(dO[...])。当 seq_len % block_M != 0 时，例如 s=130 且 auto-fit 选到 block_M=64，最后一个 tile 会从 [128, 192) 读 Q/dO，超过真实 seq_len=130；当前只给 bos/eos 的 BlockRange 做了 padding，写回 dQ 虽有 guard，但越界读已经发生，训练反向会有 illegal memory access 或未定义结果风险。现有 test_hysparse_windowed_swa.py 的 backward case 使用 s=128，没有覆盖这个尾块。
处理要求：请针对该评论修复并提交新的 commit。
建议在 kernel 内按行 guard 加载 Q/dO（或在 host 侧同步 padding q/do/lse/delta 后再 trim dq）。直接在 kernel 里改的形状如下：
for i, d in T.Parallel(BM, D):
row = bm * BM + i
in_range = row < seq_len
safe_row = T.if_then_else(in_range, row, 0)
Q_shared[i, d] = T.if_then_else(
in_range, Q[bb, safe_row, bh, d], T.cast(0, dtype)
)
for i, d in T.Parallel(BM, D_v):
row = bm * BM + i
in_range = row < seq_len
safe_row = T.if_then_else(in_range, row, 0)
dO_shared[i, d] = T.if_then_else(
in_range, dO[bb, safe_row, bh, d], T.cast(0, dtype)
)同时补一个 s % block_M != 0 的 sliding_window_mqa_attention backward 测试，例如 s=130 或显式 block_M=64 的路径，防止尾块回归。
This comment was marked as outdated.
This comment was marked as outdated.
Sorry, something went wrong.
This comment was marked as outdated.
This comment was marked as outdated.
Sorry, something went wrong.
This comment was marked as outdated.
This comment was marked as outdated.
Sorry, something went wrong.
risemeup1111
left a comment
risemeup1111
left a comment
There was a problem hiding this comment.
已复查最新提交。当前仍有阻塞问题：新提交在 transformer_layer.py 引入了 Gemma4、act offload 和 MoE 子类兼容性回归，详情见 inline comments；同时此前两个 FA4 P1 修复被回退，已在原线程补充。之前关于 HySparse recompute 顺序和 windowed backward 尾块读取的 P1 线程在当前代码中也仍未解决，沿用原线程不重复开新评论。请修复并提交新的 commit 后再复查。
| if isinstance(sublayers_spec.mlp.layer, type) and issubclass( | ||
| sublayers_spec.mlp.layer, MoELayer | ||
| ): | ||
| if sublayers_spec.mlp.layer == MoELayer: |
There was a problem hiding this comment.
这里把 MoE 判断从 issubclass(..., MoELayer) 收窄成了精确相等，导致现有 Gemma4MoELayer(MoELayer) 走到 unknown MLP 分支，additional_mlp_kwargs 不再传 pg_collection。gpt_layer_specs.py 的 attention_layer_type == "gemma4" 当前就是 LayerSpec(layer=Gemma4MoELayer, ...)，而 Gemma4MoELayer.__init__ 会把 pg_collection 继续传给基类和 router；不传会破坏 Gemma4 MoE 构建/并行配置。
处理要求：请针对该评论修复并提交新的 commit。
建议恢复子类判断：
| if sublayers_spec.mlp.layer == MoELayer: | |
| if isinstance(sublayers_spec.mlp.layer, type) and issubclass( | |
| sublayers_spec.mlp.layer, MoELayer | |
| ): |
There was a problem hiding this comment.
已在当前 head 复查，MoE 判断已恢复为 isinstance(..., type) + issubclass(..., MoELayer)，Gemma4MoELayer 这类 MoELayer 子类会继续拿到 pg_collection。这条问题我视为已修复。
| ) | ||
| return offload_kwargs | ||
|
|
||
| def build_schedule_node(self): |
There was a problem hiding this comment.
这个提交删除了 TransformerLayer._compute_act_offload_kwargs()，但仓库现有 tests/single_card_tests/ai_edited_test/transformer/test_ai_transformer_layer.py::TestDecoderlayerActOffloadSettings 仍直接调用 TransformerLayer._compute_act_offload_kwargs(fake)；当前 head 会变成 AttributeError。同时 base forward() 里也移除了 offload_kwargs = self._compute_act_offload_kwargs() 和 recompute(..., **offload_kwargs)，会让 decoderlayer_act_offload_settings 在 full/recovery recompute 路径上静默失效。
处理要求：请针对该评论修复并提交新的 commit。
这里需要同时恢复方法和 recompute 透传，形状如下：
def _compute_act_offload_kwargs(self):
decoderlayer_act_offload_settings = self.config.get(
"decoderlayer_act_offload_settings", {"type": "", "value": ""}
) or {"type": "", "value": ""}
...
return offload_kwargs
...
offload_kwargs = self._compute_act_offload_kwargs()
outputs = recompute(
self._forward_impl,
...,
**offload_kwargs,
)| if mtp_tmp_dict is not None: | ||
| rst = {**rst, **mtp_tmp_dict} | ||
| output_grad = output_grad + tuple(mtp_tmp_grad) | ||
| return rst, output_grad |
There was a problem hiding this comment.
这个 hunk 删除了 Gemma4TransformerLayerSublayersSpec 和 Gemma4TransformerLayer，但当前代码仍然引用它们：src/paddlefleet/models/gpt/gemma4_layer_specs.py 直接 from paddlefleet.transformer.transformer_layer import Gemma4TransformerLayer，gpt_layer_specs.py 在 attention_layer_type == "gemma4" 路径导入 Gemma4TransformerLayerSublayersSpec，tests/single_card_tests/transformer/test_gemma4.py 也覆盖这些符号。当前 head 下 Gemma4 layer spec 构建会在导入阶段失败。
处理要求：请针对该评论修复并提交新的 commit。
请恢复这些定义，或把所有引用和测试同步迁移到新的定义位置；需要保留的接口形状至少包括：
@dataclass
class Gemma4TransformerLayerSublayersSpec(TransformerLayerSublayersSpec):
post_self_attn_layernorm: LayerSpec | type = IdentityOp
pre_mlp_layernorm: LayerSpec | type = IdentityOp
post_mlp_layernorm: LayerSpec | type = IdentityOp
class Gemma4TransformerLayer(TransformerLayer):
def __init__(...):
...
def _forward_impl(...):
...There was a problem hiding this comment.
已在当前 head 复查，Gemma4TransformerLayerSublayersSpec 和 Gemma4TransformerLayer 已恢复到 transformer_layer.py，gemma4_layer_specs.py / gpt_layer_specs.py 的导入路径也重新可用。这条问题我视为已修复。
| hysparse_forward_mqa, | ||
| select_topk_blocks, | ||
| ) | ||
| from .block_score_fa4 import block_score_fa4_attn_fwd |
There was a problem hiding this comment.
新提交 99e25e6 又把这个问题带回来了：hysparse/__init__.py 仍执行 from .block_score_fa4 import block_score_fa4_attn_fwd，而当前 block_score_fa4.py 在模块顶层导入 paddlefleet_ops.flash_mask.cute.flashmask_utils/interface。相比上一版，_import_fa4() / is_flash_mask_available() 懒加载也被删除，所以非 SM100 或无 flash_mask 的环境会在导入 paddlefleet.tilelang_ops.hysparse 时先失败，连只使用 sliding_window_mqa_attention 的路径也无法加载。
处理要求：请针对该评论修复并提交新的 commit。
请恢复 package 级入口的懒加载/可用性检查，例如：
# hysparse/__init__.py 不要在模块初始化时导入 FA4-only 依赖
def block_score_fa4_attn_fwd(*args, **kwargs):
from .block_score_fa4 import block_score_fa4_attn_fwd as _impl
return _impl(*args, **kwargs)| ], | ||
| ) | ||
| # Grads for the tensor inputs that require grad: q, k, v only. | ||
| return dq, dk, dv |
There was a problem hiding this comment.
新提交 99e25e6 回退了这个修复，而且当前 forward 还新增了 block_bos 这个 Tensor 输入：ctx.save_for_backward(q, k, v, startend_row_indices, out, lse) 仍会在 startend_row_indices is None 时保存非 Tensor，backward 仍只 return dq, dk, dv。当前 Tensor 输入至少包含 q/k/v/block_logit/block_bos，有 mask 时还包含 startend_row_indices，因此反向返回槽位还缺 block_logit、block_bos（以及 mask Tensor 对应的 None）。
处理要求：请针对该评论修复并提交新的 commit。
建议恢复“只保存 Tensor + 显式 None 槽位”的形状，并把新增的 block_bos 也纳入返回契约：
ctx.save_for_backward(q, k, v, out, lse)
ctx.startend_row_indices = startend_row_indices
ctx.has_startend = startend_row_indices is not None
ctx.needs_grad = (
not q.stop_gradient,
not k.stop_gradient,
not v.stop_gradient,
)
...
gq, gk, gv = ctx.needs_grad
grads = [
dq if gq else None,
dk if gk else None,
dv if gv else None,
None, # block_logit
None, # block_bos
]
if ctx.has_startend:
grads.append(None) # startend_row_indices
return tuple(grads)Bring in upstream/develop's activation-offload / origin_input_ids threading in TransformerLayer (resolving the transformer_layer.py conflict) and restore the _compute_act_offload_kwargs method, then apply it to HySparseTransformerLayer's recompute path (recompute inside the recovery window, offload_kwargs threaded). Review-comment fixes: - hysparse/__init__.py: lazy-import block_score_fa4 (FA4 cute is SM10-only) so swa/select_topk import on non-Blackwell hardware. - block_score_fa4.py: PyLayer backward returns one grad slot per forward Tensor input (None for block_logit/block_bos/startend_row_indices), guards the optional startend_row_indices save, marks lse non-differentiable. - windowed_mqa_attn_bwd.py: guarded per-row Q/dO load to avoid OOB when seq_len % block_M != 0. - multi_latent_attention.py: check the HySparse full-attention branch before recompute_core_attention so block_indices is always produced. - gpt_layer_specs.py: reject enable_hy_sparse_attention combined with enable_hyper_connections / block_attention_residuals. - test_hysparse_windowed_swa.py: add ragged seqlen (S=130) backward regression test for the windowed-bwd OOB guard.
|
Rebased onto latest
Verified locally on B30Z (sm100): |
risemeup1111
left a comment
risemeup1111
left a comment
There was a problem hiding this comment.
已复查新提交。此前关于 FA4 懒加载、FA4 PyLayer 返回槽位、HySparse full-attention/recompute 顺序、windowed backward 尾块越界读、activation offload 透传的阻塞问题看起来已经修复。
仍有阻塞项需要继续处理：此前关于 MoE 子类判断和 Gemma4 Transformer layer 定义的行级线程仍未修复，请继续按原线程提交修复；另外我新增了一条行级评论，enable_mtp_magic_send 的 MTP embedding spec 被删除后会导致 MTP forward 缺少 mtp_input_embeds。
| if config.enable_hyper_connections | ||
| else None, | ||
| tail_empty_layers=tail_empty_layers_spec, | ||
| mtp=mtp_layers_spec, |
There was a problem hiding this comment.
新提交删除了 MTPEmbeddingLayer 的 spec 生成和 mtp_embedding=mtp_embedding_spec 注入。当前 GPTSublayersSpec 仍有 mtp_embedding 字段，GPTModel.get_layer_desc_list() 只有在 spec.mtp_embedding 存在时才会插入 MTPEmbeddingLayer；而 MultiTokenPredictionLayer.forward() 在 enable_mtp_magic_send=True 时会要求 dict_args["mtp_input_embeds"]，否则直接抛出 RuntimeError("MTPEmbeddingLayer may not have been executed.")。因此 enable_mtp_magic_send=True && num_nextn_predict_layers > 0 的 GPT spec 会构建出没有 MTP re-embedding 层的 pipeline，MTP forward 会失败。
处理要求：请针对该评论修复并提交新的 commit。
请恢复 magic-send 的 embedding spec 和注入，形状如下：
from paddlefleet.models.gpt.mtp_embedding_layer import MTPEmbeddingLayer
...
mtp_embedding_spec = None
if config.enable_mtp_magic_send and config.num_nextn_predict_layers > 0:
mtp_embedding_spec = LayerSpec(
layer=MTPEmbeddingLayer,
extra_kwargs={"config": config},
)
...
GPTSublayersSpec(
...
mtp=mtp_layers_spec,
mtp_embedding=mtp_embedding_spec,
mtp_lm_head=mtp_lm_head_spec,
...
)There was a problem hiding this comment.
已在当前 head 复查，MTPEmbeddingLayer 的 import、mtp_embedding_spec 构造以及 GPTSublayersSpec 的 mtp_embedding 注入都已恢复，这条问题我视为已修复。
This comment was marked as outdated.
This comment was marked as outdated.
Sorry, something went wrong.
This comment was marked as outdated.
This comment was marked as outdated.
Sorry, something went wrong.
This comment was marked as outdated.
This comment was marked as outdated.
Sorry, something went wrong.
…-count guard, test skip - transformer_layer HySparse SWA path: raise a clear ValueError with layer_number when a SWA layer receives no shared KV/block-index state from a preceding full-attention layer, instead of passing [None, None] and crashing later in MQASelfAttention with an opaque AttributeError. - block_sparse_mqa_dsa: reject H > 64 query heads up front with a clear ValueError (FlashMLA sparse fwd fixes h_q at 64 and only supports zero-padding H <= 64), instead of failing deep in the CUDA op on a mismatched h_q. - test_mqa_self_attention.test_kv_sharing: skip when the FA4 FlashMask / cuDNN DSA backends are unavailable (non-Blackwell / CPU-only), so the test no longer errors at import/op-launch time on unsupported hardware.
This comment was marked as outdated.
This comment was marked as outdated.
Sorry, something went wrong.
Wire the HySparse block-attention operators (block-score selection + block-sparse gather) into the MLA-absorbed MQA attention path, gated by the enable_hy_sparse_attention config flag. The full-attention layer runs the FA4-fused block-score kernel to score and select top-k key blocks; the SWA layer consumes the shared KV + top-k and gathers only the selected blocks. Includes the TileLang MQA/MHA block-score and block-sparse ops (fwd/bwd), the FA4-fused block-score wrapper, differentiable PyLayer wrappers, and the full op + network test suites. The FA4 fused block-score kernel is not yet merged into flash-attention (PaddlePaddle/flash-attention PR PaddlePaddle#164), so assert_fa4_block_logit_available fails fast at model-build time when HySparse is enabled but the installed FA4 build lacks the block_logit argument.
Migrate the HySparse MQA/MLA sparse-attention integration to the FA4-fused full block-score path (block_score_fa4) and the windowed MQA flash-attention kernels behind sliding-window attention, dropping the standalone TileLang block-score/MHA scoring kernels and exploratory benchmark/reference modules. Add single-card unit tests bringing the hysparse package to 95% line coverage: pipeline eq.(3) block-score + TopK selection, block-sparse gather fwd/bwd vs a dense masked reference (+ differentiable wrapper), windowed/SWA fwd/bwd, and FA4 block-score forward TopK consistency + backward gradient checks.
…ords, guards - block_score_fa4: lazy FA4 imports + is_flash_mask_available(); PyLayer now saves only tensors (stash optional mask on ctx), marks lse non-differentiable, and returns None grad slots for block_logit / startend_row_indices. - multi_latent_attention: run HySparse full-attention before the recompute branch (fixes UnboundLocalError of block_indices); remap FA4 absolute per-block logits to document-relative coords before top-k selection. - transformer_layer: raise ValueError when a SWA layer has no shared KV/top-k from a preceding full-attention layer instead of passing None. - gpt_layer_specs: require MLA attention when enable_hy_sparse_attention. - test_mqa_self_attention: skip when FA4 flash_mask (cc>=10) is unavailable. Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
The transformer-layer import chain pulls in MoELayer -> sonicmoe, which fails whenever the installed paddlefleet_ops wheel lags the develop source, making the previous HySparseTransformerLayer-based test uncollectable. Drive MLASelfAttention / MQASelfAttention directly (following the attention-level pattern of test_swa_gqa.py) so the test runs standalone while still covering the full HySparse cross-layer contract: MQA falls back to MLA when the full path is not entered, and a full layer's FA4 block-score path shares its compressed-KV latent + top-k block indices with a downstream SWA layer's block-sparse branch. A no-op TransformerLayer._log_md5 stub is injected only when the real module is unimportable. Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
Finalize the HySparse MQA/MLA sparse-attention path: * Bump the flash-attention submodule to pick up FA4-fused block-score with document-relative block bucketing (pack-equivalence) and the unified scaled block_logit contract. * block_score_fa4: thread per-query document bos (from valid_range) into the fused scorer so packed (bos>0) block selection is bit-identical to running each document alone; add a structural valid_range shape assertion. * Migrate the block-sparse gather to the DSA backend (block_sparse_mqa_dsa); drop the superseded tilelang block_sparse_attn_mqa fwd/bwd + differentiable_ops. * multi_latent_attention: build_hysparse_valid_range derives per-token [bos,eos) from the flashmask doc boundaries; add always-on structural assertions on the flashmask layout (host-side, zero training-hot-path cost). * Tests: multi-doc / packed block-score + gather pack-equivalence, whole-flow precision, DSA fwd+bwd grads, and a CPU-only build_hysparse_valid_range suite; trim redundant configs while preserving coverage.
Bring in upstream/develop's activation-offload / origin_input_ids threading in TransformerLayer (resolving the transformer_layer.py conflict) and restore the _compute_act_offload_kwargs method, then apply it to HySparseTransformerLayer's recompute path (recompute inside the recovery window, offload_kwargs threaded). Review-comment fixes: - hysparse/__init__.py: lazy-import block_score_fa4 (FA4 cute is SM10-only) so swa/select_topk import on non-Blackwell hardware. - block_score_fa4.py: PyLayer backward returns one grad slot per forward Tensor input (None for block_logit/block_bos/startend_row_indices), guards the optional startend_row_indices save, marks lse non-differentiable. - windowed_mqa_attn_bwd.py: guarded per-row Q/dO load to avoid OOB when seq_len % block_M != 0. - multi_latent_attention.py: check the HySparse full-attention branch before recompute_core_attention so block_indices is always produced. - gpt_layer_specs.py: reject enable_hy_sparse_attention combined with enable_hyper_connections / block_attention_residuals. - test_hysparse_windowed_swa.py: add ragged seqlen (S=130) backward regression test for the windowed-bwd OOB guard.
The FA4 block-score kernel changes live in the flash-attention repo and are submitted separately to its upstream. This PR should not move the paddlefleet_ops flash-attention submodule pointer; the bump will land via a dedicated submodule update once the flash-attention PR merges. Reset the gitlink back to develop's base so the PR diff carries only the PaddleFleet-side HySparse changes.
The HySparse branch was built on base PaddlePaddle#1452 via whole-tree snapshot-diff commits, which accidentally removed several unrelated features that existed in the base. Restore them so this PR contains only genuine HySparse changes: - transformer_config.py: decoderlayer_act_offload_settings, auto_subbatch_mode, moe_allgather_gate_overlap default (False -> True) - transformer_layer.py: Gemma4TransformerLayerSublayersSpec, Gemma4TransformerLayer HySparse additions (enable_hy_sparse_attention / hy_sparse_block_size / hy_sparse_topk) are preserved unchanged.
gpt_layer_specs.py was regenerated by the snapshot-diff commits and lost the MTP magic-send re-embedding wiring that existed in base PaddlePaddle#1452: - import MTPEmbeddingLayer - build mtp_embedding_spec when enable_mtp_magic_send and num_nextn_predict_layers > 0 - pass mtp_embedding=mtp_embedding_spec into GPTSublayersSpec HySparse additions (MQASelfAttention / HySparseTransformerLayer wiring and guards) are preserved unchanged.
…-count guard, test skip - transformer_layer HySparse SWA path: raise a clear ValueError with layer_number when a SWA layer receives no shared KV/block-index state from a preceding full-attention layer, instead of passing [None, None] and crashing later in MQASelfAttention with an opaque AttributeError. - block_sparse_mqa_dsa: reject H > 64 query heads up front with a clear ValueError (FlashMLA sparse fwd fixes h_q at 64 and only supports zero-padding H <= 64), instead of failing deep in the CUDA op on a mismatched h_q. - test_mqa_self_attention.test_kv_sharing: skip when the FA4 FlashMask / cuDNN DSA backends are unavailable (non-Blackwell / CPU-only), so the test no longer errors at import/op-launch time on unsupported hardware.
…ing, MoELayer issubclass check) Co-Authored-By: Claude <noreply@anthropic.com>
risemeup1111
left a comment
risemeup1111
left a comment
There was a problem hiding this comment.
已复查当前提交，Nyanpasu 此前提出的阻塞问题均已在当前 head 中修复，未发现新的需要阻塞合入的问题。
当前 Build Fleet whl 失败点是拉取构建镜像超时，不是代码编译错误；后续可重跑 CI 确认。
PaddlePaddle-bot
left a comment
PaddlePaddle-bot
left a comment
There was a problem hiding this comment.
🤖 Paddle-CI-Agent | pr_review |
2026-07-15 17:12:11
📋 Review 摘要
PR 概述：为 MLA/MQA 接入 HySparse block-score、SWA windowed MQA、DSA sparse gather，并改造 GPT layer spec/Transformer layer state 传递。
变更范围：src/paddlefleet/tilelang_ops/hysparse/、src/paddlefleet/cudnn_ops/、src/paddlefleet/transformer/、src/paddlefleet/models/gpt/、单卡测试。
影响面 Tag：Models Transformer OP
问题
| 级别 | 文件 | 概述 |
|---|---|---|
| 🔴 Bug | src/paddlefleet/models/gpt/gpt_layer_specs.py:941 |
GPTSublayersSpec 未定义 mtp_embedding，构建 GPT spec 会直接 TypeError |
| 🔴 Bug | src/paddlefleet/transformer/transformer_layer.py:1732 |
HySparse 新层未传递 origin_input_ids，MoE experimental dataflow 的 aux/z loss mask 会错 |
| 🟡 建议 | PR 级别 | 本 PR 变更量较大（27 文件 / 约 6362 行），建议按算子、模型接入、测试拆分以降低回归定位成本 |
历史 Findings 修复情况
| Finding | 问题 | 状态 |
|---|---|---|
| F1 | MoELayer.forward 传入未声明的 origin_input_ids |
✅ 已修复 |
📝 PR 规范检查
符合规范，标题已使用官方 Tag，描述包含 PR Category / PR Types / Description；历史规范建议已修复。
总体评价
建议拆分方案：
- PR 1: HySparse/DSA 算子与 PyLayer —
src/paddlefleet/tilelang_ops/hysparse/*,src/paddlefleet/cudnn_ops/* - PR 2: Transformer/GPT 接入 —
src/paddlefleet/transformer/*,src/paddlefleet/models/gpt/gpt_layer_specs.py - PR 3: 单卡测试与精度回归 —
tests/single_card_tests/**/test_hysparse_*,tests/single_card_tests/test_mqa_self_attention.py
本轮按风险优先审查了 GPT spec 构建、HySparseTransformerLayer shared state/MoE 传参、MQA/DSA PyLayer 契约、SWA kernel 测试；未全量覆盖所有 TileLang 数值分支和 benchmark 删除后的外部引用。
| else None, | ||
| tail_empty_layers=tail_empty_layers_spec, | ||
| mtp=mtp_layers_spec, | ||
| mtp_embedding=mtp_embedding_spec, |
There was a problem hiding this comment.
🔴 Bug 这里给 GPTSublayersSpec 传入了未声明的 mtp_embedding 关键字。
当前 src/paddlefleet/models/gpt/gpt_model.py 里的 GPTSublayersSpec dataclass 只有 embedding/head_empty_layers/.../mtp_loss 等字段，没有 mtp_embedding。因此 get_gpt_spec() 构造 GPTSublayersSpec(...) 时会直接抛 TypeError: unexpected keyword argument 'mtp_embedding'，即使 enable_mtp_magic_send=False 也会传入 None 并影响所有 GPT spec 构建。
建议修复方式：要么把 mtp_embedding 字段补到 GPTSublayersSpec 并在 GPTModel.get_layer_desc_list() 中真正插入对应 layer；要么在完整 wiring 落地前删除这个关键字传参。
| ) | ||
| self._log_md5(hidden_states, "post_attn_residual", self.layer_number) | ||
| with profile(timer_name): | ||
| output = self._forward_mlp(hidden_states, input_ids=input_ids) |
There was a problem hiding this comment.
🔴 Bug HySparse 新层调用 _forward_mlp 时丢掉了 origin_input_ids。
基础 TransformerLayer 会从 dict_args 取 origin_input_ids 并传给 _forward_mlp，随后 MoELayer.forward/MoERouter 用它在 gpt_model_use_experimental_version=True 时计算 aux/z loss 的原始 token mask。这里新增的 HySparseTransformerLayer._forward_impl 既没有接收 origin_input_ids，也没有在 recompute 分支传入，最后只调用 _forward_mlp(hidden_states, input_ids=input_ids)；HySparse + MoE + experimental dataflow 下会退回使用截断/局部 input_ids，padding/MTP token 的分母会错。
建议修复方式：和基类保持一致，在 forward() 的 recompute 参数、_forward_impl 签名中传递 origin_input_ids=dict_args.get("origin_input_ids")，并改为 self._forward_mlp(hidden_states, input_ids=input_ids, origin_input_ids=origin_input_ids)。
PR Category
Performance Optimization
PR Types
New features
Description
Integrate HySparse block-sparse attention (arXiv:2602.03560), MLA-absorbed MQA
variant, into the model:
tilelang_ops/hysparse): MQA/MHA block-score, MQA block-sparsegather, sliding-window MQA, and the FA4-fused full block-score kernel
(
block_score_fa4), withpaddle.autograd.PyLayerdifferentiable wrappersand a dense reference.
MQASelfAttention/HySparseTransformerLayer/gpt_layer_specsconnect the full-attention layer (produces shared compressed KV + top-k block
indices) to the SWA layer (consumes them for the block-sparse branch), gated
by
enable_hy_sparse_attention.tests.
Review fixes in this update:
is_flash_mask_available(); save onlytensors, stash the optional mask on
ctx, marklsenon-differentiable, andreturn
Nonegrad slots forblock_logit/startend_row_indices.UnboundLocalErrorofblock_indices).top-k so packed multi-document selection aligns with the sparse gather.
ValueErrorwhen a SWA layer runs without a preceding full-attentionlayer's shared state; require MLA attention when enabling HySparse.
unavailable (e.g. Hopper H20 CI runners).