Skip to content

Add ReplaySSM support for linear attention - #1585

Draft
sufubao wants to merge 1 commit into
ModelTC:mainfrom
sufubao:support-replayssm
Draft

sufubao wants to merge 1 commit into
ModelTC:mainfrom
sufubao:support-replayssm

Conversation

@sufubao

@sufubao sufubao commented Sep 20, 2026

Copy link
Copy Markdown
Collaborator

问题与结果

Qwen3-Next/Qwen3.5 的线性注意力在 speculative verify 时会为每个候选 token 保存完整 SSM state,状态写入和显存占用随 MTP 宽度增长。本 PR 增加默认关闭的 --enable_replayssm:GDN FP32 使用单份 checkpoint 和接受历史,GDN BF16/KDA MTP 使用紧凑 raw-input cache,避免保存候选状态快照。

实现

  • GDN FP32 verify 只重建输出;接受后移动请求游标,历史满时才更新 checkpoint。
  • checkpoint flush、CPU cache 和 PD 导出从 raw k/v 与 FP32 gate/beta 按原生递推重放,派生 delta 不进入长期状态。
  • BF16 与 KDA 使用每轮提交的 Compact cache,保留逐 token 舍入语义;普通 BF16/KDA decode 继续使用现有 recurrent kernel。
  • prefix cache、PD 传输和请求复用前统一 materialize/reset;功能通过新参数显式开启。
  • KDA 路径为 GLM-5.3-Flash 预留,模型接入仍由 feat: add GLM-5.3 Flash with MTP and vision support #1575 负责,本 PR 不在 main 上开放 GLM。

验证

  • m39 H200:34 passed, 4 skipped,覆盖 accepted prefix、不同窗口、CUDA Graph、循环 flush、CPU cache、PD page、请求复用以及 Compact/native 对比。实验记录:260920-170629-*
  • 48 层、width=4、接受 3 token 的内核参考形状:batch 32/68/96 相对原生分别为 1.06x / 1.09x / 1.16x;Replay state/history 约占原生 speculative SSM snapshot 的 44%。实验记录:260920-165530-*
  • SGLang 风格最小 16-token 矩阵 verify 在当前形状上明显回退,因此未保留;实现使用顺序 output-only verify。

未覆盖与评审重点

  • 当前数据是内核测试,尚未在目标 SenseNova 权重上证明端到端 OTPS 收益;因此保持 opt-in。
  • FP32 verify 会因浮点重结合出现生成差异,不承诺 token 级 bitwise parity。
  • 请重点检查 accepted cursor、自动 flush 与 CPU/PD materialize 的状态边界,以及 KDA Compact 对 feat: add GLM-5.3 Flash with MTP and vision support #1575 的适配接口。

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