ModelTC · sufubao · Sep 20, 2026
diff --git a/docs/replayssm/README.md b/docs/replayssm/README.md
@@ -0,0 +1,151 @@
+# ReplaySSM implementation and validation
+
+ReplaySSM is opt-in: add `--enable_replayssm` to an existing Qwen3-Next or
+Qwen3.5 command. `--replayssm_cache_len` accepts 16 (default), 32, or 64 and must
+cover `mtp_step + 1`. Existing scheduling, convolution, cache-page formats,
+prefill backends, and acceptance policy remain the owners of their current work.
+
+## Runtime paths
+
+| Recurrent state | Ordinary decode | Speculative verification |
+| --- | --- | --- |
+| GDN FP32 | Deferred checkpoint + accepted update history | Same representation; commit advances accepted prefix |
+| GDN BF16 | Existing recurrent kernel | Compact inputs; replay accepted prefix with per-token rounding |
+| GLM KDA, after model integration | Existing recurrent kernel | Compact inputs with bounded per-key gate |
+
+FP32 history stores normalized keys, deltas, and cumulative log decay. Two
+alternating history regions prevent cross-CTA overwrite races during folding.
+Metadata is shared across layers; HOLD requests never acquire history. The verify
+kernel processes candidates sequentially without writing full candidate states.
+It does **not** implement SGLang's parallel triangular-solve verify formulation.
+
+Before prefill reentry, CPU checkpoint saves, or PD state export, pending accepted
+history is folded into canonical state. Restore/reuse clears history metadata.
+CPU and PD snapshots select the accepted convolution window independently of the
+recurrent cursor. Compact commit already leaves canonical recurrent state.
+
+BF16 uses compact recomputation because deferring rounded state updates changes
+the recurrence. Ordinary BF16/KDA decode keeps the existing kernel; compact replay
+is useful for avoiding speculative full-state snapshots, not as a blanket
+replacement for ordinary recurrence.
+
+## GLM-5.3-Flash
+
+The compact cache already supports KDA's per-key gate, but model routing is left
+to [PR #1575](https://github.com/ModelTC/LightLLM/pull/1575). The two changes were
+tested together at `7eef17bb360c9b454ed5707ea1175cafed8446dd`: kernel/MTP/cache/PD
+tests and fixed/dynamic KDA adapter tests passed. Full GLM serving and multi-node
+PD remain unvalidated, so this change does not advertise GLM support on main.
+
+## Exact fold on m39
+
+The GDN path separates output reconstruction from checkpoint maintenance.
+Normalized keys and deltas are used only for verify outputs; flush and cache/PD
+materialization replay accepted raw BF16 keys/values plus FP32 gates/betas into
+the FP32 checkpoint. Approximate output records therefore never become the next
+persistent state.
+
+Verify uses the sequential output-only kernel. A minimum 16-token matrix tile
+was rejected: it measured 0.15--0.17x of native at width 4 and 0.36--0.47x at
+width 16 on the 48-layer reference shape. On an H200 in m39, width 4 with three
+accepted tokens measured 1.06x/1.09x/1.16x at batches 32/68/96. Replay state and
+history occupied about 44% of the native speculative SSM snapshot allocation.
+Sources: `260920-165210-*` (rejected matrix-only result) and `260920-165530-*`
+(sequential result); the width-16 matrix rejection is `260920-170004-*`.
+Sequential width 16 improved to 0.79x but still regressed (`260920-170210-*`).
+These are kernel measurements, not end-to-end serving gains.
+
+## Measurements on m33
+
+Hardware: H200; PyTorch 2.11.0+cu130, Triton 3.6. All measurements are recorded with
+`exp -m`; IDs below identify `/root/experiments/runs` inside `lightllm-replayssm`
+and exported copies under the m33 user's `~/experiments/runs`.
+
+Kernel results include metadata, accepted-prefix commit, and recurring folds.
+Shape: K=V=128, 16 query heads, 32 value heads, batch 256, capacity 16.
+
+| Path | Verify width | Baseline ms | Replay ms | Speedup |
+| --- | ---: | ---: | ---: | ---: |
+| FP32 deferred | 1 | 0.35152 | 0.24927 | 1.41x |
+| FP32 deferred | 3 | 0.85617 | 0.51800 | 1.65x |
+
+Source: `260918-093117-*` (BV64/warps2).
+Earlier BF16 compact measurements (`260918-094913-*`, 2.64–3.51x) used BV32.
+Repeated comparisons against the actual native BV8 default exposed a rounding
+difference. Final GDN compact now uses BV8 and matches native output/state over
+32 acceptance cycles in the test. **The earlier compact timings are superseded**;
+Final compact was remeasured on an idle H200 on September 20. With width3,
+all candidates accepted, and one layer with eight value heads, speedups at
+batches 1/4/16/32/68 were 0.90/0.95/0.98/1.12/1.28x (`260920-013240-*`).
+A 48-layer chain with one final cross-layer commit, four key heads and twelve
+value heads per layer, measured 1.54/1.28/1.31/1.50/1.56x at those batches
+(`260920-013503-*`). The latter uses the local public Qwen3.5-27B configuration
+divided by TP4 as a reference shape; it is not full-model serving. Both use CUDA
+Graph timing and omit projections, full attention, communication, draft, sampling,
+scheduling, and PD transfer. Single-layer timing does not predict the multilayer
+chain, and neither establishes an end-to-end gain.
+At width4 with three accepted tokens, the same 48-layer shape measured
+3.512/2.023 ms (native/compact) at batch32, 7.514/4.007 ms at batch68, and
+10.770/5.588 ms at batch96 (`260920-015547-*`). These H200 measurements support
+testing MTP3 on the H100 deployment, but do not predict its service OTPS.
+The speculative baseline uses the existing default BV8/warps1 configuration,
+not an exhaustive baseline autotune. Small batches can regress: FP32 batch1 was
+0.42x for ordinary decode, 0.49x for width3. These are kernel results, not service
+speedups or proof of peak performance. Opt-in remains appropriate.
+
+At width3/batch256, FP32 state plus history was 809,505,796 bytes versus
+1,616,904,192 bytes baseline. BF16 compact used 282,313,472 versus 808,452,096 bytes.
+These totals exclude unrelated model/KV memory.
+
+Qwen3.5-0.8B serving exercised ordinary decode, dynamic MTP, and repeated prefix
+cache hits (512 cached tokens). Four of five deterministic prompts matched all
+128 tokens. One diverged after token57 for deferred ordinary decode and token61
+for dynamic MTP, so bitwise generation parity is **not** claimed. Reassociation
+changes numerical results even when recurrence tests pass.
+
+The existing GSM8K script, first100 test questions after five few-shot examples,
+scored 31/100 for baseline, FP32 deferred ReplaySSM, and BF16 compact dynamic MTP. Sources:
+`260918-095130-*`, `260919-141406-*`, and `260919-141716-*`. This is a small accuracy smoke test.
+The matching native BF16 dynamic-MTP baseline scored 32/100
+(`260919-141956-*`). Several answers differed even in serial repeated requests
+(`260919-142125-*`), prompting the BV8 correction above; initial BF16 serving
+results describe the superseded BV32 kernel. The later runs shared GPUs with
+other workloads; their latency is not comparable.
+The corrected BV8 compact run scored 30/100 (`260919-142538-*`), and four of
+five serial probes matched native text (`260919-142626-*`). One probe still
+differed, so full serving numerical equivalence remains under investigation.
+Repeating the native concurrent baseline scored 33/100 (`260919-142927-*`),
+showing that its 100-question score itself varies with execution conditions.
+A diagnostic service ran native verification alongside compact GDN on identical
+live inputs, then checked accepted state after commit. Across ten requests on
+these five probes it logged no output/state mismatches (`260919-143150-*` client
+run; diagnostic server source/log archived with the experiment). This localizes
+the remaining service differences outside the tested recurrence operations;
+it does not establish whole-model numerical parity.
+
+## Reproduce
+
+Run on a CUDA machine from the repository root:
+
+```sh
+python -m pytest -q unit_tests/common/basemodel/triton_kernel/linear_att/test_replayssm.py
+PYTHONPATH=. exp -m 'ReplaySSM deferred kernels' python test/benchmark/kernels/benchmark_replayssm.py
+PYTHONPATH=. exp -m 'ReplaySSM compact kernels' python test/benchmark/kernels/benchmark_replayssm.py --compact --state-dtype bfloat16 --widths 3 5
+PYTHONPATH=. exp -m 'ReplaySSM serving parity' python test/benchmark/service/compare_replayssm.py --baseline-port 18932 --replay-port 18934
+```
+
+The main runtime suite passed 65 tests (two native KDA comparisons skipped on
+main); all four extended native comparisons passed with the PR kernel injected.
+Tests cover accepted prefixes, irregular verify groups, graph replay, alternating
+ring folds, materialization, request reuse, CPU checkpoints, PD page helpers, and
+compact/native recurrence comparisons. KDA native comparisons skip on main and
+run against the PR's real kernel when it is present. The initial single-window KDA case matched BF16 exactly and had FP32 state error
+at most 7.45e-9. Expanded changing-input tests exposed rounding differences across
+repeated commits; KDA therefore uses explicit FP32/BF16 error tolerances and does
+not promise bitwise equivalence. GDN compact retains exact native comparison.
+
+Remaining qualification: larger task accuracy samples, isolated end-to-end
+throughput, tuned-baseline comparisons, full GLM serving, and multi-node PD.
+
+Algorithm references: [ReplaySSM](https://dao-lab.ai/blog/2026/replayssm/) and
+[SGLang](https://github.com/sgl-project/sglang/tree/bbfcda48cebfbde2ed4be4e8273eb593c764e638).
diff --git a/lightllm/common/basemodel/attention/linear/gdn.py b/lightllm/common/basemodel/attention/linear/gdn.py
@@ -35,7 +35,6 @@ def get_prefill_kernel(self):
         pass
 
     def _init_linear_layer_metadata(self, network_config, tp_world_size):
-
         self.mtp_step = get_env_start_args().mtp_step
 
         # Linear attention specific dimensions
@@ -107,18 +106,17 @@ def create_att_decode_state(self, infer_state: "InferStateInfo") -> "LinearAttDe
 
 @dataclasses.dataclass
 class LinearAttPrefillAttState(BasePrefillAttState):
-
     b_conv_buffer_idx: torch.Tensor = None
     b_ssm_buffer_idx: torch.Tensor = None
 
     def init_state(self):
-        backend: LinearAttBackend = self.backend
-        mtp_step = backend.mtp_step
         # 每次 _prefill 都会在 runtime infer_state 上调用 init_state。
         # prefill cuda graph 回调必须走 new_infer_state.prefill_att_state1，
         # 才能读到这里按当前 batch（含 token padding 后的 dummy request）更新的索引。
         self.b_conv_buffer_idx = self.infer_state.b_req_idx
-        self.b_ssm_buffer_idx = self.infer_state.b_req_idx * (mtp_step + 1)
+        self.b_ssm_buffer_idx = self.infer_state.b_req_idx * self.infer_state.req_manager.ssm_slots_per_req
+        if self.infer_state.req_manager.replay_cache is not None:
+            self.infer_state.req_manager.replay_cache.materialize(self.infer_state.b_req_idx)
         return
 
     def prefill_att(
@@ -199,12 +197,13 @@ def _gdn_prefill_kernel(
 
 @dataclasses.dataclass
 class LinearAttDecodeAttState(BaseDecodeAttState):
-
     b_conv_buffer_idx: torch.Tensor = None
     b_ssm_buffer_idx: torch.Tensor = None
     b1_mtp_cu_q_seq_len: torch.Tensor = None
     b_num_accepted_tokens: torch.Tensor = None
 
+    b_replay_positions: torch.Tensor = None
+
     def init_state(self):
         draft_step = self.backend.model.mtp_manager.get_decode_draft_step(self.backend.model.is_mtp_draft_model)
         if draft_step == 0:
@@ -213,6 +212,11 @@ def init_state(self):
             self._init_dynamic_mtp_decode_state(draft_step + 1)
         else:
             self._init_fixed_mtp_decode_state(draft_step)
+        replay = getattr(self.infer_state.req_manager, "replay_cache", None)
+        if replay is not None:
+            self.b_replay_positions = replay.positions(self.b_conv_buffer_idx)
+            if draft_step == 0:
+                replay.commit(self.b_conv_buffer_idx)
 
     def _init_normal_decode_state(self):
         self.b_conv_buffer_idx = self.infer_state.b_req_idx
@@ -252,6 +256,9 @@ def _init_fixed_mtp_decode_state(self, draft_step: int):
         self._init_mtp_ssm_buffer_idx(mtp_size)
 
     def _init_mtp_ssm_buffer_idx(self, mtp_size: int):
+        if getattr(self.infer_state.req_manager, "replay_cache", None) is not None:
+            self.b_ssm_buffer_idx = self.b_conv_buffer_idx
+            return
         att_batch_size = self.b_conv_buffer_idx.shape[0]
         # Each request owns mtp_size consecutive recurrent-state slots.
         b_ssm_buffer_start_idx = (self.b_conv_buffer_idx * mtp_size).view(att_batch_size, 1)
@@ -339,6 +346,27 @@ def _gdn_decode_kernel(
             backend.tp_num_v_heads,
             backend.head_v_dim,
         )
+        replay = infer_state.req_manager.replay_cache
+        if replay is not None:
+            layer = (
+                layer_weight.layer_num_
+                - layer_weight.layer_num_ // infer_state.req_manager.linear_config.full_attention_interval
+            )
+            return (
+                replay.forward(
+                    layer,
+                    query,
+                    key,
+                    value,
+                    a,
+                    b,
+                    layer_weight.linear_A_log.weight,
+                    layer_weight.linear_dt_bias.weight,
+                    self.b_conv_buffer_idx,
+                    self.b_replay_positions,
+                ),
+                z,
+            )
         core_attn_out, _ = fused_recurrent_gated_delta_rule(
             q=query,
             k=key,
@@ -384,6 +412,25 @@ def _gdn_mtp_kernel(
         )
 
         query, key, value = backend._rearrange_mixed_qkv(mixed_qkv, decode=False)
+        replay = infer_state.req_manager.replay_cache
+        if replay is not None:
+            layer = (
+                layer_weight.layer_num_
+                - layer_weight.layer_num_ // infer_state.req_manager.linear_config.full_attention_interval
+            )
+            return replay.forward(
+                layer,
+                query,
+                key,
+                value,
+                a,
+                b,
+                layer_weight.linear_A_log.weight,
+                layer_weight.linear_dt_bias.weight,
+                self.b_conv_buffer_idx,
+                self.b_replay_positions,
+                cu_seqlens_q,
+            )
         assert self.b_ssm_buffer_idx.dim() == 2, "SSM buffer idx must be 2D [N, S+1]"
         # #8b: b_num_accepted_tokens >= 1 is guaranteed upstream: init/cache restore set 1,
         # and MTP decode only writes values in [1, mtp_step+1]. The old per-layer per-step