Skip to content

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

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

sufubao wants to merge 23 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 保存本层参数与共享 backend,通过 attention.context(infer_state) 取得当前 context,再显式调用 forward。context 按 allocate_context → prepare → update → forward 执行;State 继续负责存储、缓存恢复和释放。保留 group_id。无持久 KV 的视觉/音频路径使用 VarlenAttentionOp,保留 --vit_att_backend。

本轮收敛了共享 backend 的状态归属:MTP cursor 放入层参数,GDN/KDA native/replay 共用 decode layout,按 State 分别取得 accepted count;backend 不再浅拷贝注入 spec。context 使用 slots dataclass,workspace requirement 和 replay State 接口有明确类型;BatchBuilder 集中借用 InferState 的输入 tensor。selector 显式接收 specs、设备能力和执行约束,不再反向查询 basemodel。Dense/MLA FA3 共用 decode planning。out 明确由调用方拥有,支持 provider 返回 padded/strided buffer 的有效 view,也允许临时结果复制;不承诺零分配。import 按标准库、第三方、lightllm、相对导入分组,模型调用及中英文文档同步更新。

已合并 upstream/main 9640f76,包含 ReplaySSM #1690、KDA prefill 融合 #1705、对角求解并行化 #1703 和 decode 文件句柄修复 #1707。新 KDA kernel 除 import 外计算 AST 与上游一致。

本机 H200 验证(实验命令、源码哈希和失败记录在 /data6/exp_logs/runs):

  • attention/provider/resource/Graph、GLM5、FP8、varlen 和 speculative 回归去重后 2421 passed,0 skipped;首轮2358通过后因旧 monkeypatch 指向已删除的 MLA import 停止,修正夹具后相关模块与剩余用例97通过(261112-220026/220342)。
  • ReplaySSM acceptance/Graph、动态 MTP 与 KDA 融合专项122通过(215807);随后合并最新上游,新增求解器逐位对拍、KDA/GLM和启动路径专项286通过(220525)。以上测试有重叠,不相加。
  • Qwen3.5-0.8B 完整权重,BF16、TP1/DP1,native/replay MTP step1各64次主模型/draft logits比较,最大误差0(215414/215428);随后仅移除dense helper未用参数,最终普通replay decode32次比较误差0(220037)。覆盖prefill/decode Graph、GPU/CPU页恢复、本地DP/PD buffer恢复和分配器回收。
  • PR的330个Python文件通过Black21.12b0 / Flake8 6.1.0。

之前扩展测试中的旧native MTP test_mtp_verify_path[1]仍有torch.equal最大差值0.000977;合并前测试/kernel覆盖对照可复现(180025),未放宽断言,本轮不计为通过。此前ReplaySSM完整624例结果见历史实验,本轮重跑的是上述受影响路径。跨节点PD、完整多模态模型、全部MTP组合及性能未验收。

仍保留共享AttentionBatch字段名和现有候选表;没有引入统一Registry或强制所有provider直接写out。Optional字段用于尚未prepare的context,泛型尚未经过全仓静态类型检查。无关工作区修改未纳入提交。

@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