[RunningLeon]

Summary

• Add opt-in ReplaySSM for supported Gated DeltaNet/KDA inference, including MTP accepted-input commits.
• Fix early-stop continuation by invalidating private KV/SSM state at prefill admission and recomputing retained history.
• Correct ARSpec MRoPE positions after resetting the recompute step.

Performance test settings

Setting	Value
Models / parallelism	Qwen3.5-35B-A3B: TP1; GLM5.3 Flash: TP4
Backends	LMDeploy and SGLang
Matrix	Target-only and MTP3, each with Replay OFF/ON; 16 arms total
Requests	1,024 measured requests per arm; 4 warmup requests excluded
Input / output	16,384 input tokens / 2,048 output tokens; fixed-length generation, EOS ignored
Target-only	Maximum batch size 256; client concurrency 256
MTP3	3 draft tokens; maximum batch size 64; client concurrency 64
Replay buffer	L16 for all Replay ON arms
Graph / prefix caching	CUDA Graph ON; prefix caching OFF
Prefill / context limits	Maximum prefill tokens 8,192; context/session length 32,768
SSM state	FP32
Qwen KV capacity	Automatic allocation
GLM KV capacity	Fixed: LMDeploy 32,768 blocks × 64 tokens; SGLang max_total_num_tokens=2,097,152, page_size=64

Fixed capacity aligns nominal token capacity, not physical cache bytes across backends. Automatic allocation may convert state-memory savings into additional KV capacity.

Performance results

Each row compares Replay ON against OFF within the same backend, model and mode. MTP=OFF denotes target-only. Throughput is output tokens/s; TTFT and TPOT are means. Memory is the measured whole-server NVML peak per GPU; TP4 uses the mean of the four per-GPU peaks. Negative memory deltas indicate savings.

Backend	Model	Blocks	MTP	Batch	Replay L	Output tok/s Δ ↑	TTFT mean Δ ↓	TPOT mean Δ ↓	Peak NVML Δ (GiB/GPU) ↓
LMDeploy	Qwen3.5 TP1	Automatic	OFF	256	16	-2.49%	+4.79%	-2.25%	+1.152 (+0.95%)
LMDeploy	Qwen3.5 TP1	Automatic	MTP3	64	16	+1.63%	-5.19%	-1.56%	-3.041 (-2.48%)
SGLang	Qwen3.5 TP1	Automatic	OFF	256	16	+1.30%	-1.01%	-1.77%	+2.840 (+2.21%)
SGLang	Qwen3.5 TP1	Automatic	MTP3	64	16	+7.94%	+0.35%	-7.64%	+1.904 (+1.46%)
LMDeploy	GLM5.3 Flash TP4	Fixed	OFF	256	16	+4.39%	-4.72%	-3.25%	+4.618 (+4.17%)
LMDeploy	GLM5.3 Flash TP4	Fixed	MTP3	64	16	+0.17%	+0.62%	-0.19%	-4.396 (-3.79%)
SGLang	GLM5.3 Flash TP4	Fixed	OFF	256	16	-1.09%	+0.29%	+1.06%	+2.992 (+2.60%)
SGLang	GLM5.3 Flash TP4	Fixed	MTP3	64	16	+1.11%	+0.99%	-1.23%	+0.000 (+0.00%)

Conclusions

• LMDeploy Qwen MTP: saves 3.041 GiB/GPU (2.48%), with throughput +1.63%, TTFT -5.19%, and TPOT -1.56%.
• LMDeploy GLM MTP: saves 4.396 GiB/GPU (3.79%). Performance is nearly unchanged: throughput +0.17%, TTFT +0.62%, and TPOT -0.19%. Memory reduction is the primary benefit.
• LMDeploy target-only: increases memory use. Qwen automatic capacity regresses in throughput and TTFT; GLM fixed capacity improves both. Results do not support a universal speedup claim.
• SGLang MTP: improves throughput and TPOT, but slightly increases TTFT and shows no whole-server memory reduction. Cache-allocation differences limit direct cross-backend comparisons.

BC-breaking (Optional)

Does the modification introduce changes that break the backward-compatibility of the downstream repositories?
If so, please describe how it breaks the compatibility and how the downstream projects should modify their code to keep compatibility with this PR.

Use cases (Optional)

If this PR introduces a new feature, it is better to list some use cases here, and update the documentation.

Checklist

1. Pre-commit or other linting tools are used to fix the potential lint issues.
2. The modification is covered by complete unit tests. If not, please add more unit tests to ensure the correctness.
3. If the modification has a dependency on downstream projects of a newer version, this PR should be tested with all supported versions of downstream projects.
4. The documentation has been modified accordingly, like docstring or example tutorials.

[Copilot]

🟡 Changes recommended

MTP workspaces scale with persistent prefix-cache slots rather than active batch capacity, risking excessive GPU memory use or OOM.

3 open findings

• Size workspace capacity by active batch instead of total state slots · New
• Correct value-head state shape from [HV,V,K] to [V,K] · New
• Correct value-head state shape from [HV,V,K] to [V,K] · New

What changed in this PR

Adds opt-in ReplaySSM support for standard Gated DeltaNet and Qwen3.5 MTP decoding, including cache lifecycle, CUDA graphs, configuration, kernels, tests, and documentation.

Changes:

• Adds standard and MTP ReplaySSM Triton kernels and state-cache schemas.
• Integrates cursor commits, checkpoint flushing, CUDA graphs, configuration, and CLI options.
• Adds GPU/CPU regression tests and bilingual documentation.

File	Description
tests/​pytorch/​spec_decode/​test_cudagraph_strategy.py	Tests graph-stable flush buffers.
tests/​pytorch/​nn/​test_replayssm_flag_off.py	Tests disabled-path parity.
tests/​pytorch/​kernel/​test_gated_delta_replayssm.py	Tests standard kernels and lifecycle.
tests/​pytorch/​kernel/​test_gated_delta_replayssm_spec.py	Tests MTP reconstruction and commits.
tests/​pytorch/​engine/​test_model_agent.py	Tests agent integration and checkpoint handling.
tests/​pytorch/​engine/​test_inputs_maker.py	Tests flush-mask generation.
tests/​pytorch/​engine/​test_cache_inputs.py	Tests mask transfer and stream recording.
tests/​pytorch/​config/​test_replayssm_config.py	Tests configuration and cache schemas.
lmdeploy/​pytorch/​nn/​gated_delta.py	Defines ReplaySSM state views.
lmdeploy/​pytorch/​models/​utils/​cudagraph.py	Adds fixed graph flush buffers.
lmdeploy/​pytorch/​models/​qwen3_next.py	Integrates ReplaySSM state into Qwen3-Next.
lmdeploy/​pytorch/​models/​qwen3_5.py	Integrates ReplaySSM state into Qwen3.5.
lmdeploy/​pytorch/​model_inputs.py	Extends model and step contexts.
lmdeploy/​pytorch/​kernels/​cuda/​gated_delta_replayssm.py	Implements standard ReplaySSM kernels.
lmdeploy/​pytorch/​kernels/​cuda/​gated_delta_replayssm_spec.py	Implements MTP ReplaySSM kernels.
lmdeploy/​pytorch/​kernels/​cuda/​gated_delta_replayssm_config.py	Defines fixed kernel launch configurations.
lmdeploy/​pytorch/​kernels/​cuda/​gated_delta_preprocess.py	Supports compact key-head preprocessing.
lmdeploy/​pytorch/​engine/​model_agent/​agent.py	Manages flush and MTP commit lifecycle.
lmdeploy/​pytorch/​engine/​inputs_maker.py	Creates checkpoint flush masks.
lmdeploy/​pytorch/​engine/​executor/​base.py	Adds ReplaySSM memory accounting.
lmdeploy/​pytorch/​engine/​executor/​__init__.py	Validates supported configurations.
lmdeploy/​pytorch/​engine/​config_builder.py	Propagates public settings.
lmdeploy/​pytorch/​engine/​cache_inputs.py	Transfers flush masks to devices.
lmdeploy/​pytorch/​consts.py	Defines named cache keys.
lmdeploy/​pytorch/​configurations/​qwen3_next.py	Builds ReplaySSM cache schemas.
lmdeploy/​pytorch/​configurations/​qwen3_5.py	Enables schemas for Qwen3.5.
lmdeploy/​pytorch/​config.py	Adds ReplaySSM configuration fields.
lmdeploy/​pytorch/​backends/​gated_delta_rule.py	Extends preprocessing metadata and API.
lmdeploy/​pytorch/​backends/​cuda/​graph_runner/​full_graph.py	Marks standard ReplaySSM graphs.
lmdeploy/​pytorch/​backends/​cuda/​gated_delta_rule.py	Dispatches ReplaySSM implementations.
lmdeploy/​messages.py	Exposes public engine options.
lmdeploy/​cli/​utils.py	Defines ReplaySSM CLI arguments.
lmdeploy/​cli/​serve.py	Connects options to API serving.
lmdeploy/​cli/​cli.py	Exposes options for interactive chat.
docs/​zh_cn/​index.rst	Links Chinese documentation.
docs/​zh_cn/​advance/​replayssm.md	Documents ReplaySSM in Chinese.
docs/​en/​index.rst	Links English documentation.
docs/​en/​advance/​replayssm.md	Documents ReplaySSM in English.

🧠 Review effort: Balanced

---

💡 Add a code-review agent skill or configure MCP servers for context-aware, tailored reviews. Learn more in the docs.

[caikun-pjlab]

LGTM