Skip to content

Add ReplaySSM state modes for GDN and KDA - #1585

Draft
sufubao wants to merge 10 commits into
ModelTC:mainfrom
sufubao:support-replayssm
Draft

sufubao wants to merge 10 commits into
ModelTC:mainfrom
sufubao:support-replayssm

Conversation

@sufubao

@sufubao sufubao commented Sep 20, 2026 •

Copy link
Copy Markdown
Collaborator

为 GDN 和 GLM-5.3-FLASH KDA 增加可选的 compact / replay 状态存储,默认 native。KDA 按 key 维衰减,inline/precompute 均使用正向衰减的低秩历史重建;历史 key 递推更新,避免每 token 重算整段指数。

复用请求接受、fold、CPU checkpoint 和 PD 生命周期,接入真实 prefill、普通 decode、固定/动态 MTP,并保留 GLM indexer tail。调参只用独立 scratch,区分 recurrence/lower_bound;按 token 数和 sequence 容量分别选配置,在各 CUDA Graph 捕获前固定;模型输出携带 verify 配置,去 padding 和双 microbatch 合并 accept 后仍使用对应布局。Replay checkpoint 使用固定布局,后续 batch 调优不会改变已捕获 Graph。缩短 KDA 调参缓存文件名,支持 TP1/2/4/8 的持久化。实际模型还复现并修复了 native Graph 的大批量 empty/HOLD DeepGEMM 调度问题:初始化零值 dummy pool,真实长度继续屏蔽输出;独立 paged/nonpaged logits 与严格 scratch 越界检查通过。

  • 前序完整验证 H100:实际 GLM 镜像 Torch2.13/Triton3.7.1 946 passed、0 skipped;数学内核相同的 Torch2.11/Triton3.6 版本 891 passed、0 skipped。覆盖独立 FP32/BF16 递推、L4–64、全部27候选、部分接受、Graph/HOLD、活动 checkpoint、跨 TP/模式 PD。
  • FP32 native/compact 对照保持零容差。BF16 比较容纳实测的编译器舍入传播,负向实验仍拒绝缺失逐 token 舍入的实现。Replay 的舍入边界不同,不保证 BF16 生成文本逐 token 相同。
  • 本次按 batch 配置改动:H100 / Torch2.11.0+cu130 / Triton3.6.0 775 passed、0 skipped。包含不同配置的 Graph 交替 replay、固定/动态 padding、双 microbatch 合并 accept、GDN/KDA、FP32/BF16,以及 replay inline/precompute;输出和接受后状态按相同配置逐元素一致。实验由 exp 记录。
  • Pinned Black/flake8、whitespace 通过。

GLM 实际34层KDA,TP8/BF16 state/MTP2:native每请求25.5 MiB/卡,replay L8 inline约17.02、precompute约14.89 MiB,含历史/游标,不含其他缓存。没有MTP时 replay 比单 native checkpoint 多占历史空间。

同 H100 的34层完整更新周期(prepare/forward/accept/fold/export/merge,5次重复,8轮,27候选):BF16 precompute MTP2 在batch64/192更新时间减少约8.5–24.8%,batch16增加约16–29%。这是状态更新组件结果,排除conv、模型/MoE与HTTP;不承诺端到端吞吐收益。

实际GLM纯文本服务(TP8/BF16state/MTP2dynamic/L8,同源码与固定容量):四profile各1050测量请求,共4200,无HTTP失败。Replay inline相对native吞吐中位−1.31%至+3.06%,TPOT−1.48%至+6.64%;precompute吞吐−2.94%至+0.62%。短TTFT波动、FP8/BF16生成差异及高draft接受率的synthetic负载限制外推,不声明普遍端到端加速或正式任务质量无回退。当前推荐inline作为GLM验证配置;raw token/时间和4条自然前缀对照完整保存。视觉峰值探测的现成镜像ABI问题另存,使用既有--disable_vision --disable_audio完成文本对照。

@sufubao sufubao changed the title Add ReplaySSM support for linear attention Add ReplaySSM state modes for GDN Oct 6, 2026
@sufubao sufubao changed the title Add ReplaySSM state modes for GDN Add ReplaySSM state modes for GDN and KDA Oct 6, 2026
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