Skip to content

refactor(attention): separate backends, runtime and persistent state - #1700

Draft
sufubao wants to merge 21 commits into
ModelTC:mainfrom
sufubao:refactor/attention-backend-plans
Draft

sufubao wants to merge 21 commits into
ModelTC:mainfrom
sufubao:refactor/attention-backend-plans

Conversation

@sufubao

@sufubao sufubao commented Nov 10, 2026 •

Copy link
Copy Markdown
Collaborator

将 attention 计算、批次执行上下文和 KV/Conv/SSM 持久状态分开。模型按 LayerSpec 分组,LayerAttentionBinding 保存本层参数与共享 AttentionBackend,显式执行 attention.forward(context, inputs, out=...)。AttentionExecutionContext 通过 allocate_context → prepare → update → forward 管理批次 metadata 和 workspace;持久状态、缓存恢复和释放仍由 State 管理。保留 group_id。

视觉/音频完整 Q/K/V 路径使用 VarlenAttentionOp,统一 FA3/Triton/SDPA/xFormers 的选择和调用;保留 --vit_att_backend。能力检查通过时返回 None,不支持时抛 UnsupportedAttention。模型、Graph、测试及中英文接入文档同步迁移,import 按标准库、第三方、lightllm、相对导入分组。

已 merge upstream/main dc53088,包含 ReplaySSM #1690 和 KDA prefill 融合 #1705。GDN/KDA replay 迁移到新 backend/context 接口,保留上游 kernel、history/fold 和状态转移语义;补齐 State 布局声明与 Graph 预热快照,MTP 从输出绑定的 commit 计算每个请求的验证宽度。FlashQla replay 仅支持 FP16/BF16 prefill。GLM5 FP8 固定 gather/compaction workspace 在 KV 容量计算前预留,每个 group/microbatch 独立;Graph replay 刷新 KV 行映射,小 gather buffer 走分块合并。

当前 ad6abb9,本机 H200 验证:

  • 最新上游 KDA prefill 融合迁移至 native KDA backend;新 kernel 与上游字节一致。新增卷积/QK norm/history 逐位对拍、可变边界 Graph、backend 集成及现有 KDA/MTP/layer/kernel 回归 177 passed、0 skipped(181811)。

  • attention/provider/resource/Graph、GLM5 数值、MTP、循环层映射与启动参数回归 2498 passed,0 failed/skipped(261112-180130)。随后能力限制和 workspace 定向复测 147 passed(180342),小 gather buffer 分块及 FP8 全组 120 passed(180606)。

  • ReplaySSM 新增 kernel/state/acceptance/Graph/精度矩阵:624 个不同用例最终全部通过、0 skipped。首轮 553 passed 后遇到上游 padding 夹具的同步行数不足;按 bucket 修正测试夹具后,相关组及剩余用例 83 passed(175817/181345),未修改生产 kernel 或放宽数值检查。

  • GLM5 FP8、缓存/PD、共享请求管理和 ReplaySSM 快照检查 592 passed(175613),与上述矩阵有重叠,不累计为独立用例数。

  • Qwen3.5-0.8B 完整权重,BF16、TP1/DP1、FA3 dense + Triton GDN,native/replay 两种 SSM 模式均通过 MTP step=1、prefill/decode Graph、GPU/CPU 页恢复和本地 DP/PD buffer 恢复;每种模式 64 次主模型/draft logits 比较最大误差 0,分配器回收通过(175245/175315)。关闭 MTP 的 replay 普通 decode 也通过相同 Graph/缓存恢复检查,32 次 logits 比较误差 0(181032)。后续仅收紧 FlashQla replay 能力限制和修改 GLM5 路径,所测 Triton GDN forward 未变。

  • PR 相对新上游的 326 个 Python 文件通过 Black 21.12b0 / Flake8 6.1.0;7 个迁移的 ReplaySSM kernel/decoder 除 import 外计算 AST 与上游一致。

扩展运行旧 native MTP kernel 测试时,test_mtp_verify_path[1] 的 torch.equal 断言失败,最大差值 0.000977。用合并前 4c79bac 的测试与两个 kernel 文件覆盖对照,复现相同失败(180025);共享依赖仍为合并工作树。未修改断言或放宽容差,此项不计入通过结果。

命令、源码哈希、失败日志及结果保存在 /data6/exp_logs/runs。依赖为 torch 2.11.0+cu130、sglang-kernel 0.4.5、FlashMLA 5dcead、FlashInfer 0.6.18 JIT。结果覆盖所列本地配置;跨节点 PD、完整多模态模型、全部 MTP 组合及性能未验收。此前 HTTP/varlen 探针结果仍见历史记录,本轮未重跑 HTTP 服务。无关工作区修改未纳入提交。

@sufubao
sufubao force-pushed the refactor/attention-backend-plans branch 3 times, most recently from 09fab80 to 011b909 Compare November 10, 2026 15:12
@sufubao sufubao changed the title refactor(attention): separate backend plans from persistent state refactor(attention): unify text and encoder backend interfaces Nov 10, 2026
@sufubao
sufubao force-pushed the refactor/attention-backend-plans branch from 011b909 to 592cbca Compare November 11, 2026 00:44
@sufubao sufubao changed the title refactor(attention): unify text and encoder backend interfaces refactor(attention): separate backends, runtime and persistent state Nov 11, 2026
@sufubao
sufubao force-pushed the refactor/attention-backend-plans branch from 690e372 to ab254cd Compare November 11, 2026 02:55
@sufubao
sufubao force-pushed the refactor/attention-backend-plans branch from 167964f to b9618da Compare November 12, 2026 05:24
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant