From 762c235848121e12d13cf6a7562e31a0bed711b5 Mon Sep 17 00:00:00 2001 From: sufubao Date: Sun, 20 Sep 2026 17:06:24 +0800 Subject: [PATCH] feat: add ReplaySSM support for linear attention --- docs/replayssm/README.md | 151 +++++++++ .../common/basemodel/attention/linear/gdn.py | 59 +++- .../triton_kernel/linear_att/replayssm.py | 303 +++++++++++++++++ .../linear_att/replayssm_compact.py | 225 ++++++++++++ .../triton_kernel/linear_att_copy.py | 10 + .../qwen3next_mem_manager.py | 16 +- lightllm/common/req_manager/linear_att.py | 68 +++- lightllm/server/api_cli.py | 13 + lightllm/server/api_start.py | 12 + lightllm/server/core/objs/start_args_type.py | 2 + test/benchmark/kernels/benchmark_replayssm.py | 141 ++++++++ test/benchmark/service/compare_replayssm.py | 77 +++++ .../linear_att/test_replayssm.py | 320 ++++++++++++++++++ 13 files changed, 1382 insertions(+), 15 deletions(-) create mode 100644 docs/replayssm/README.md create mode 100644 lightllm/common/basemodel/triton_kernel/linear_att/replayssm.py create mode 100644 lightllm/common/basemodel/triton_kernel/linear_att/replayssm_compact.py create mode 100644 test/benchmark/kernels/benchmark_replayssm.py create mode 100644 test/benchmark/service/compare_replayssm.py create mode 100644 unit_tests/common/basemodel/triton_kernel/linear_att/test_replayssm.py diff --git a/docs/replayssm/README.md b/docs/replayssm/README.md new file mode 100644 index 0000000000..c04633b498 --- /dev/null +++ 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 index ca6ceaec43..fc1acd3242 100644 --- 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 diff --git a/lightllm/common/basemodel/triton_kernel/linear_att/replayssm.py b/lightllm/common/basemodel/triton_kernel/linear_att/replayssm.py new file mode 100644 index 0000000000..f12a674db6 --- /dev/null +++ b/lightllm/common/basemodel/triton_kernel/linear_att/replayssm.py @@ -0,0 +1,303 @@ +# SPDX-License-Identifier: Apache-2.0 +"""GDN ReplaySSM in LightLLM's [K, V] state layout. + +The checkpoint contains only folded tokens. Derived records reconstruct verify +outputs, while raw inputs replay the accepted suffix into the checkpoint. Verify +may overwrite the uncommitted suffix; only acceptance advances the cursor. +Algorithm reference: https://dao-lab.ai/blog/2026/replayssm/ +""" + +import torch +import triton +import triton.language as tl + + +@triton.jit +def _positions( + cursors, + reqs, + out, + accepted, + N: tl.constexpr, + HOLD: tl.constexpr, + L: tl.constexpr, + WIDTH: tl.constexpr, + COMMIT: tl.constexpr, + BLOCK: tl.constexpr, +): + i = tl.program_id(0) * BLOCK + tl.arange(0, BLOCK) + r = tl.load(reqs + i, i < N, HOLD) + valid = (i < N) & (r != HOLD) + cursor = tl.load(cursors + r, valid, 0) + n = cursor % (2 * L) + phase = cursor & (2 * L) + if COMMIT: + count = tl.load(accepted + r, valid, 0) + 1 if WIDTH > 1 else 1 + next_cursor = tl.where(n + WIDTH > L, phase ^ (2 * L), phase + n) + count + tl.store(cursors + r, next_cursor, valid) + else: + tl.store(out + i, cursor, i < N) + + + + +@triton.jit +def _replay( + Q, Kp, Vp, A, B, Alog, Bias, State, Keys, Deltas, Gates, RawKeys, RawValues, Betas, + Reqs, Cursors, Cu, Out, + SQ: tl.constexpr, SK: tl.constexpr, SV: tl.constexpr, SA: tl.constexpr, SB: tl.constexpr, + H: tl.constexpr, HV: tl.constexpr, K: tl.constexpr, V: tl.constexpr, + L: tl.constexpr, WIDTH: tl.constexpr, HOLD: tl.constexpr, + BK: tl.constexpr, BV: tl.constexpr, VARLEN: tl.constexpr, +): + """Sequential output-only path; faster than a padded MMA tile for short MTP windows.""" + iv, seq, hv = tl.program_id(0), tl.program_id(1), tl.program_id(2) + req = tl.load(Reqs + seq).to(tl.int64) + if VARLEN: + start, end = tl.load(Cu + seq), tl.load(Cu + seq + 1) + else: + start, end = seq, seq + 1 + vv = iv * BV + tl.arange(0, BV) + if req == HOLD: + for t in range(start, end): + tl.store(Out + (t * HV + hv) * V + vv, 0, vv < V) + return + if start == end: + return + + kk = tl.arange(0, BK) + ll = tl.arange(0, L) + slot = req * HV + hv + cursor = tl.load(Cursors + seq) + n = cursor % (2 * L) + base = (cursor // (2 * L)) * L + sp = State + slot * K * V + kk[:, None] * V + vv[None, :] + state = tl.load(sp, (kk[:, None] < K) & (vv[None, :] < V), 0).to(tl.float32) + + if n + WIDTH > L: + for j in range(n): + raw_k = tl.load(RawKeys + (slot * (2 * L) + base + j) * K + kk, kk < K, 0).to(tl.float32) + raw_v = tl.load(RawValues + (slot * (2 * L) + base + j) * V + vv, vv < V, 0).to(tl.float32) + g = tl.load(Gates + slot * (2 * L) + base + j).to(tl.float32) + beta = tl.load(Betas + slot * (2 * L) + base + j).to(tl.float32) + raw_k /= tl.sqrt(tl.sum(raw_k * raw_k) + 1.0e-6) + state *= tl.exp(g) + d = beta * (raw_v - tl.sum(state * raw_k[:, None], 0)) + state += raw_k[:, None] * d[None, :] + tl.store(sp, state, (kk[:, None] < K) & (vv[None, :] < V)) + base, n = L - base, 0 + + history = ll < n + keys = tl.load( + Keys + (slot * (2 * L) + base) * K + ll[:, None] * K + kk[None, :], + history[:, None] & (kk[None, :] < K), 0, + ).to(tl.float32) + deltas = tl.load( + Deltas + (slot * (2 * L) + base) * V + ll[:, None] * V + vv[None, :], + history[:, None] & (vv[None, :] < V), 0, + ).to(tl.float32) + gates = tl.load(Gates + slot * (2 * L) + base + ll, history, 0).to(tl.float32) + gate_prefix = tl.cumsum(gates, axis=0) + total_g = tl.sum(gates, axis=0) + h = hv // (HV // H) + log_a, bias = tl.load(Alog + hv).to(tl.float32), tl.load(Bias + hv).to(tl.float32) + + for t in range(start, end): + q = tl.load(Q + t * SQ + h * K + kk, kk < K, 0).to(tl.float32) + raw_k = tl.load(Kp + t * SK + h * K + kk, kk < K, 0).to(tl.float32) + v = tl.load(Vp + t * SV + hv * V + vv, vv < V, 0).to(tl.float32) + q = q / tl.sqrt(tl.sum(q * q) + 1.0e-6) * (K**-0.5) + k = raw_k / tl.sqrt(tl.sum(raw_k * raw_k) + 1.0e-6) + x = tl.load(A + t * SA + hv).to(tl.float32) + bias + g = -tl.exp(log_a) * tl.where(x <= 20.0, tl.log(1.0 + tl.exp(x)), x) + beta = tl.sigmoid(tl.load(B + t * SB + hv).to(tl.float32)) + total_g += g + weights = tl.where(history, tl.exp(total_g - gate_prefix), 0.0) + hk = tl.sum(keys * k[None, :], axis=1) * weights + hq = tl.sum(keys * q[None, :], axis=1) * weights + sk = tl.sum(state * k[:, None], axis=0) * tl.exp(total_g) + tl.sum(deltas * hk[:, None], axis=0) + sq = tl.sum(state * q[:, None], axis=0) * tl.exp(total_g) + tl.sum(deltas * hq[:, None], axis=0) + d = beta * (v - sk) + tl.store(Out + (t * HV + hv) * V + vv, sq + d * tl.sum(k * q), vv < V) + record = base + n + tl.store(Deltas + (slot * (2 * L) + record) * V + vv, d, vv < V) + tl.store(RawValues + (slot * (2 * L) + record) * V + vv, v, vv < V) + if iv == 0: + tl.store(Keys + (slot * (2 * L) + record) * K + kk, k, kk < K) + tl.store(RawKeys + (slot * (2 * L) + record) * K + kk, raw_k, kk < K) + tl.store(Gates + slot * (2 * L) + record, g) + tl.store(Betas + slot * (2 * L) + record, beta) + keys = tl.where((ll == n)[:, None], k[None, :], keys) + deltas = tl.where((ll == n)[:, None], d[None, :], deltas) + gate_prefix = tl.where(ll >= n, total_g, gate_prefix) + history |= ll == n + n += 1 + + +@triton.jit +def _materialize( + State, + RawKeys, + RawValues, + Gates, + Betas, + Cursors, + Reqs, + SLOTS: tl.constexpr, + HV: tl.constexpr, + K: tl.constexpr, + V: tl.constexpr, + L: tl.constexpr, + HOLD: tl.constexpr, + BK: tl.constexpr, + BV: tl.constexpr, +): + iv, row, lh = tl.program_id(0), tl.program_id(1), tl.program_id(2) + req = tl.load(Reqs + row).to(tl.int64) + if req == HOLD: + return + cursor = tl.load(Cursors + req) + n = cursor % (2 * L) + base = (cursor // (2 * L)) * L + if n == 0: + return + layer, head = lh // HV, lh % HV + slot = (layer * SLOTS + req) * HV + head + kk = tl.arange(0, BK) + vv = iv * BV + tl.arange(0, BV) + sp = State + slot * K * V + kk[:, None] * V + vv[None, :] + state = tl.load(sp, (kk[:, None] < K) & (vv[None, :] < V), 0).to(tl.float32) + for j in range(n): + raw_k = tl.load(RawKeys + (slot * (2 * L) + base + j) * K + kk, kk < K, 0).to(tl.float32) + raw_v = tl.load(RawValues + (slot * (2 * L) + base + j) * V + vv, vv < V, 0).to(tl.float32) + g = tl.load(Gates + slot * (2 * L) + base + j).to(tl.float32) + beta = tl.load(Betas + slot * (2 * L) + base + j).to(tl.float32) + raw_k /= tl.sqrt(tl.sum(raw_k * raw_k) + 1.0e-6) + state *= tl.exp(g) + d = beta * (raw_v - tl.sum(state * raw_k[:, None], 0)) + state += raw_k[:, None] * d[None, :] + tl.store(sp, state, (kk[:, None] < K) & (vv[None, :] < V)) + + +class ReplaySSMCache: + """Request-owned scratch; canonical CPU/PD checkpoints never include it.""" + + def __init__(self, state, capacity, verify_width, activation_dtype=torch.bfloat16): + assert state.dtype == torch.float32, "deferred ReplaySSM requires FP32 SSM state" + assert capacity >= max(16, verify_width) and capacity & (capacity - 1) == 0 + self.state = state + self.capacity = capacity + self.verify_width = verify_width + layers, slots, hv, k, v = state.shape + self.hold = slots - 1 + self.cursors = torch.zeros(slots, dtype=torch.int32, device=state.device) + # Alternate halves on fold so CTAs cannot overwrite history another + # V tile is still reading. Cursor packs phase (2*L) and count (0..L). + self.keys = torch.empty((layers, slots, hv, 2 * capacity, k), dtype=torch.float32, device=state.device) + self.deltas = torch.empty((layers, slots, hv, 2 * capacity, v), dtype=torch.float32, device=state.device) + self.gates = torch.empty((layers, slots, hv, 2 * capacity), dtype=torch.float32, device=state.device) + self.raw_keys = torch.empty((layers, slots, hv, 2 * capacity, k), dtype=activation_dtype, device=state.device) + self.raw_values = torch.empty( + (layers, slots, hv, 2 * capacity, v), dtype=activation_dtype, device=state.device + ) + self.betas = torch.empty((layers, slots, hv, 2 * capacity), dtype=torch.float32, device=state.device) + + def reset(self, req): + self.cursors[req] = 0 + + def positions(self, reqs): + positions = torch.empty_like(reqs) + _positions[(triton.cdiv(reqs.numel(), 256),)]( + self.cursors, + reqs, + positions, + None, + reqs.numel(), + self.hold, + self.capacity, + self.verify_width, + False, + 256, + ) + return positions + + def commit(self, reqs, accepted=None): + _positions[(triton.cdiv(reqs.numel(), 256),)]( + self.cursors, + reqs, + None, + accepted, + reqs.numel(), + self.hold, + self.capacity, + self.verify_width, + True, + 256, + ) + + def materialize(self, reqs): + if reqs.numel() == 0: + return + layers, slots, hv, k, v = self.state.shape + _materialize[(triton.cdiv(v, 32), reqs.numel(), layers * hv)]( + self.state, + self.raw_keys, + self.raw_values, + self.gates, + self.betas, + self.cursors, + reqs, + slots, + hv, + k, + v, + self.capacity, + self.hold, + triton.next_power_of_2(k), + 32, + ) + self.cursors[reqs] = 0 + + def forward(self, layer, q, k, v, a, b, a_log, bias, reqs, positions, cu_seqlens=None, run_config=None): + hv, kd, vd = self.state.shape[-3:] + axis = 1 if cu_seqlens is not None else 0 + out = torch.empty_like(v) + config = run_config or {"BV": 64, "num_warps": 2} + bv = config["BV"] + _replay[(triton.cdiv(vd, bv), reqs.numel(), hv)]( + q, + k, + v, + a, + b, + a_log, + bias, + self.state[layer], + self.keys[layer], + self.deltas[layer], + self.gates[layer], + self.raw_keys[layer], + self.raw_values[layer], + self.betas[layer], + reqs, + positions, + cu_seqlens, + out, + q.stride(axis), + k.stride(axis), + v.stride(axis), + a.stride(0), + b.stride(0), + q.shape[-2], + hv, + kd, + vd, + self.capacity, + self.verify_width, + self.hold, + triton.next_power_of_2(kd), + bv, + cu_seqlens is not None, + num_warps=config["num_warps"], + ) + return out diff --git a/lightllm/common/basemodel/triton_kernel/linear_att/replayssm_compact.py b/lightllm/common/basemodel/triton_kernel/linear_att/replayssm_compact.py new file mode 100644 index 0000000000..83ef18ad0f --- /dev/null +++ b/lightllm/common/basemodel/triton_kernel/linear_att/replayssm_compact.py @@ -0,0 +1,225 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Compact speculative state for GDN/BF16 and KDA. + +Verify retains raw inputs rather than full SSM snapshots. Commit re-executes the +accepted prefix with the same per-token rounding as the recurrent baseline. +KDA follows GLM-5.3-Flash's bounded per-key gate and rsqrt normalization. +""" + +import torch +import triton +import triton.language as tl + + +@triton.jit +def _compact( + Q, + Kp, + Vp, + A, + B, + Alog, + Bias, + State, + Keys, + Values, + Decays, + Betas, + Reqs, + Cu, + Accepted, + Out, + SQ: tl.constexpr, + SK: tl.constexpr, + SV: tl.constexpr, + SA: tl.constexpr, + SB: tl.constexpr, + SLOTS: tl.constexpr, + H: tl.constexpr, + HV: tl.constexpr, + K: tl.constexpr, + V: tl.constexpr, + WIDTH: tl.constexpr, + HOLD: tl.constexpr, + BK: tl.constexpr, + BV: tl.constexpr, + COMMIT: tl.constexpr, + KDA: tl.constexpr, + LOWER: tl.constexpr, +): + iv, row, lh = tl.program_id(0), tl.program_id(1), tl.program_id(2) + req = tl.load(Reqs + row).to(tl.int64) + layer, hv = lh // HV, lh % HV + if COMMIT: + start = 0 + end = tl.load(Accepted + req) + 1 if req != HOLD else 0 + else: + start, end = tl.load(Cu + row), tl.load(Cu + row + 1) + kk = tl.arange(0, BK) + vv = iv * BV + tl.arange(0, BV) + if req == HOLD or end == start: + if not COMMIT: + for t in range(start, end): + tl.store(Out + (t * HV + hv) * V + vv, 0, vv < V) + return + slot = (layer * SLOTS + req) * HV + hv + sp = State + slot * K * V + kk[:, None] * V + vv[None, :] + state = tl.load(sp, (kk[:, None] < K) & (vv[None, :] < V), 0).to(tl.float32) + for t in range(start, end): + record = slot * WIDTH + t - start + if COMMIT: + k = tl.load(Keys + record * K + kk, kk < K, 0).to(tl.float32) + v = tl.load(Values + record * V + vv, vv < V, 0).to(tl.float32) + decay = tl.load(Decays + record * K + kk, kk < K, 0) if KDA else tl.load(Decays + record) + beta = tl.load(Betas + record) + else: + h = hv // (HV // H) + k = tl.load(Kp + t * SK + h * K + kk, kk < K, 0).to(tl.float32) + v = tl.load(Vp + t * SV + hv * V + vv, vv < V, 0).to(tl.float32) + log_a = tl.load(Alog + hv).to(tl.float32) + if KDA: + bias = tl.load(Bias + hv * K + kk, kk < K, 0).to(tl.float32) + gate = tl.load(A + t * SA + hv * K + kk, kk < K, 0).to(tl.float32) + decay = tl.exp(LOWER * tl.sigmoid(tl.exp(log_a) * (gate + bias))) + else: + x = tl.load(A + t * SA + hv).to(tl.float32) + tl.load(Bias + hv).to(tl.float32) + g = -tl.exp(log_a) * tl.where(x <= 20.0, tl.log(1.0 + tl.exp(x)), x) + decay = tl.exp(g) + beta = tl.sigmoid(tl.load(B + t * SB + hv).to(tl.float32)) + tl.store(Values + record * V + vv, v, vv < V) + if iv == 0: + tl.store(Keys + record * K + kk, k, kk < K) + tl.store(Betas + record, beta) + if KDA: + tl.store(Decays + record * K + kk, decay, kk < K) + else: + tl.store(Decays + record, decay) + if KDA: + k = k * tl.rsqrt(tl.sum(k * k) + 1.0e-6) + state *= decay[:, None] + else: + k = k / tl.sqrt(tl.sum(k * k) + 1.0e-6) + state *= decay + d = (v - tl.sum(state * k[:, None], 0)) * beta + state += k[:, None] * d[None, :] + if not COMMIT: + q = tl.load(Q + t * SQ + h * K + kk, kk < K, 0).to(tl.float32) + if KDA: + q = q * (tl.rsqrt(tl.sum(q * q) + 1.0e-6) * (K**-0.5)) + else: + q = q / tl.sqrt(tl.sum(q * q) + 1.0e-6) * (K**-0.5) + tl.store(Out + (t * HV + hv) * V + vv, tl.sum(state * q[:, None], 0), vv < V) + state = state.to(State.dtype.element_ty).to(tl.float32) + if COMMIT: + tl.store(sp, state, (kk[:, None] < K) & (vv[None, :] < V)) + + +class CompactSSMCache: + def __init__(self, state, verify_width, activation_dtype, kind="gdn", lower_bound=-5.0): + assert kind in ("gdn", "kda") + self.state = state + self.verify_width = verify_width + self.kind = kind + self.lower_bound = lower_bound + layers, slots, hv, k, v = state.shape + self.hold = slots - 1 + shape = (layers, slots, hv, verify_width) + self.keys = torch.empty((*shape, k), device=state.device, dtype=activation_dtype) + self.values = torch.empty((*shape, v), device=state.device, dtype=activation_dtype) + self.decays = torch.empty((*shape, k) if kind == "kda" else shape, device=state.device, dtype=torch.float32) + self.betas = torch.empty(shape, device=state.device, dtype=torch.float32) + + def reset(self, req): + # Scratch is overwritten by verify before commit; it has no accepted history. + pass + + def positions(self, reqs): + return None + + def materialize(self, reqs): + # Commit always leaves a canonical checkpoint. + pass + + def forward(self, layer, q, k, v, a, b, a_log, bias, reqs, positions, cu_seqlens=None): + assert cu_seqlens is not None, "compact replay is only used for speculative verify" + _, slots, hv, kd, vd = self.state.shape + out = torch.empty_like(v) + # Match native reduction layouts: changing BV can alter BF16 rounding. + bv = 32 if self.kind == "kda" else 8 + _compact[(triton.cdiv(vd, bv), reqs.numel(), hv)]( + q, + k, + v, + a, + b, + a_log, + bias, + self.state[layer], + self.keys[layer], + self.values[layer], + self.decays[layer], + self.betas[layer], + reqs, + cu_seqlens, + None, + out, + q.stride(1), + k.stride(1), + v.stride(1), + a.stride(0), + b.stride(0), + slots, + q.shape[-2], + hv, + kd, + vd, + self.verify_width, + self.hold, + triton.next_power_of_2(kd), + bv, + False, + self.kind == "kda", + self.lower_bound, + num_warps=4 if self.kind == "kda" else 1, + ) + return out + + def commit(self, reqs, accepted): + layers, slots, hv, kd, vd = self.state.shape + bv = 32 if self.kind == "kda" else 8 + _compact[(triton.cdiv(vd, bv), reqs.numel(), layers * hv)]( + None, + None, + None, + None, + None, + None, + None, + self.state, + self.keys, + self.values, + self.decays, + self.betas, + reqs, + None, + accepted, + None, + 0, + 0, + 0, + 0, + 0, + slots, + hv, + hv, + kd, + vd, + self.verify_width, + self.hold, + triton.next_power_of_2(kd), + bv, + True, + self.kind == "kda", + self.lower_bound, + num_warps=4 if self.kind == "kda" else 1, + ) diff --git a/lightllm/common/basemodel/triton_kernel/linear_att_copy.py b/lightllm/common/basemodel/triton_kernel/linear_att_copy.py index eb17507fb0..3b60908c49 100644 --- a/lightllm/common/basemodel/triton_kernel/linear_att_copy.py +++ b/lightllm/common/basemodel/triton_kernel/linear_att_copy.py @@ -25,6 +25,9 @@ def _copy_linear_att_state_to_kv_buffer( cpu_kv_ssm_stride_l, cpu_kv_ssm_stride_d, mtp_step, + conv_offsets, + conv_element_bytes: tl.constexpr, + HAS_CONV_OFFSETS: tl.constexpr, gpu_conv_dim, # number of conv rows gpu_conv_tail_dim_bytes, # bytes copied per conv row; equals the CPU/cache row width gpu_ssm_tail_dim, @@ -53,6 +56,8 @@ def _copy_linear_att_state_to_kv_buffer( cur_state_req_idx = (cur_req_idx * (mtp_step + 1)).to(tl.int64) gpu_conv_base = gpu_conv_ptr + cur_layer * gpu_conv_stride_l + cur_req_idx * gpu_conv_stride_s + if HAS_CONV_OFFSETS: + gpu_conv_base += tl.load(conv_offsets + cur_req_idx) * conv_element_bytes cpu_conv_base = cpu_kv_conv_ptr + big_page_buffer_idx * cpu_kv_conv_stride_s + cur_layer * cpu_kv_conv_stride_l conv_tail_dim = gpu_conv_dim * gpu_conv_tail_dim_bytes for i in range(tl.cdiv(conv_tail_dim, BLOCK)): @@ -86,10 +91,12 @@ def copy_linear_att_state_to_kv_buffer( cpu_kv_conv_state: torch.Tensor, # [buffer_num, linear_layer_num, conv_dim, kernel_size] cpu_kv_ssm_state: torch.Tensor, # [buffer_num, linear_layer_num, ...] mtp_step: int, + conv_offsets: torch.Tensor = None, ): # gpu_conv_state 的后两维可能是不连续的。 assert len(b_req_idx) == big_page_buffer_ids.shape[0] BLOCK = 4096 + conv_element_bytes = gpu_conv_state.element_size() assert gpu_conv_state.dim() == 4, "gpu_conv_state must be [layer, s, conv_dim, widened_width]" assert cpu_kv_conv_state.dim() == 4, "cpu_kv_conv_state must be [size, layer, conv_dim, width_narrow]" @@ -144,6 +151,9 @@ def copy_linear_att_state_to_kv_buffer( cpu_kv_ssm_stride_l=cpu_kv_ssm_state.stride(1), cpu_kv_ssm_stride_d=cpu_kv_ssm_state.stride(2), mtp_step=mtp_step, + conv_offsets=conv_offsets, + conv_element_bytes=conv_element_bytes, + HAS_CONV_OFFSETS=conv_offsets is not None, gpu_conv_dim=gpu_conv_dim, gpu_conv_tail_dim_bytes=gpu_conv_tail_dim_bytes, gpu_ssm_tail_dim=gpu_ssm_tail_dim, diff --git a/lightllm/common/kv_cache_mem_manager/qwen3next_mem_manager.py b/lightllm/common/kv_cache_mem_manager/qwen3next_mem_manager.py index ce22808e16..0fa0df26a5 100644 --- a/lightllm/common/kv_cache_mem_manager/qwen3next_mem_manager.py +++ b/lightllm/common/kv_cache_mem_manager/qwen3next_mem_manager.py @@ -71,6 +71,9 @@ def _free_linear_att_buffers(self): return def write_to_shm(self, req_manager): + self.replay_cache = req_manager.replay_cache + self.ssm_slots_per_req = req_manager.ssm_slots_per_req + self.req_to_mtp_state_index = req_manager.req_to_mtp_state_index self.req_to_conv_state = req_manager.req_to_conv_state self.req_to_ssm_state = req_manager.req_to_ssm_state # super().write_to_shm() 会用 ForkingPickler 序列化本对象，torch 在 dump 时会把 @@ -260,7 +263,7 @@ def read_page_to_req( return def _get_req_state_indexes(self, req_idx: int): - mtp_size = get_env_start_args().mtp_step + 1 + mtp_size = getattr(self.mem_manager, "ssm_slots_per_req", get_env_start_args().mtp_step + 1) # Conv is one widened slot per request; SSM keeps the historical S+1 block layout. return req_idx, req_idx * mtp_size @@ -272,8 +275,14 @@ def _write_one_rank( conv_page: torch.Tensor, ssm_page: torch.Tensor, ): + replay = getattr(mem, "replay_cache", None) + if replay is not None: + replay.materialize(torch.tensor([req_idx], dtype=torch.int32, device=mem.req_to_ssm_state.buffer.device)) conv_req_idx, ssm_req_idx = self._get_req_state_indexes(req_idx) conv_state = mem.req_to_conv_state.buffer[:, conv_req_idx, ..., : self.conv_shape[-1]] + if replay is not None and mem.req_to_mtp_state_index is not None: + offsets = torch.arange(self.conv_shape[-1], device=conv_state.device) + mem.req_to_mtp_state_index[req_idx] + conv_state = mem.req_to_conv_state.buffer[:, conv_req_idx].index_select(-1, offsets) ssm_state = mem.req_to_ssm_state.buffer[:, ssm_req_idx, ...] self._copy_conv_state_to_page(conv_state, conv_page, mem, tp_index) self._copy_ssm_state_to_page(ssm_state, ssm_page, mem, tp_index) @@ -452,6 +461,11 @@ def _read_one_rank( conv_page: torch.Tensor, ssm_page: torch.Tensor, ): + replay = getattr(mem, "replay_cache", None) + if replay is not None: + replay.reset(req_idx) + if mem.req_to_mtp_state_index is not None: + mem.req_to_mtp_state_index[req_idx] = 0 conv_req_idx, ssm_req_idx = self._get_req_state_indexes(req_idx) conv_state = mem.req_to_conv_state.buffer[:, conv_req_idx, ..., : self.conv_shape[-1]] ssm_state = mem.req_to_ssm_state.buffer[:, ssm_req_idx, ...] diff --git a/lightllm/common/req_manager/linear_att.py b/lightllm/common/req_manager/linear_att.py index ccab8a06f6..f21c81d1d0 100644 --- a/lightllm/common/req_manager/linear_att.py +++ b/lightllm/common/req_manager/linear_att.py @@ -13,9 +13,29 @@ class ReqManagerForMamba(HybridAttentionReqManager): - def __init__(self, max_request_num, max_sequence_length, mem_manager, linear_config: LinearAttCacheConfig): + def __init__( + self, + max_request_num, + max_sequence_length, + mem_manager, + linear_config: LinearAttCacheConfig, + recurrent_kind="gdn", + ): super().__init__(max_request_num, max_sequence_length, mem_manager) - self.mtp_step = get_env_start_args().mtp_step + args = get_env_start_args() + self.mtp_step = args.mtp_step + replay = getattr(args, "enable_replayssm", False) + if ( + replay + and self.mtp_step == 0 + and (recurrent_kind == "kda" or linear_config.ssm_state_dtype != torch.float32) + ): + from lightllm.utils.log_utils import init_logger + + init_logger(__name__).info("ReplaySSM: non-speculative BF16/KDA decode uses the recurrent kernel") + replay = False + self.ssm_slots_per_req = 1 if replay else self.mtp_step + 1 + self.replay_cache = None # 因为在mtp的推理中，需要标记每个请求对应的mtp index状态(conv state 和 ssm state)，在mtp对应序列中 # 的真实位置，所以需要需要一个标记来记录，不然算子无法找到真实的处理起点。 self.req_to_mtp_state_index = ( @@ -40,21 +60,39 @@ def __init__(self, max_request_num, max_sequence_length, mem_manager, linear_con device="cuda", ) self.req_to_ssm_state = LayerCache( - size=(max_request_num + 1) * (self.mtp_step + 1), + size=(max_request_num + 1) * self.ssm_slots_per_req, dtype=self.linear_config.ssm_state_dtype, shape=self.linear_config.get_ssm_state_shape(), layer_num=self.linear_config.linear_layer_num, device="cuda", ) + if replay: + from lightllm.common.basemodel.triton_kernel.linear_att.replayssm import ReplaySSMCache + + if recurrent_kind == "gdn" and linear_config.ssm_state_dtype == torch.float32: + self.replay_cache = ReplaySSMCache( + self.req_to_ssm_state.buffer, + args.replayssm_cache_len, + self.mtp_step + 1, + linear_config.conv_state_dtype, + ) + else: + from lightllm.common.basemodel.triton_kernel.linear_att.replayssm_compact import CompactSSMCache + + self.replay_cache = CompactSSMCache( + self.req_to_ssm_state.buffer, self.mtp_step + 1, linear_config.conv_state_dtype, kind=recurrent_kind + ) return def init_hybrid_attention_state(self, req: "InferReq"): + if self.replay_cache is not None: + self.replay_cache.reset(req.req_idx) conv_index = req.req_idx - ssm_start = req.req_idx * (self.mtp_step + 1) + ssm_start = req.req_idx * self.ssm_slots_per_req self.req_to_conv_state.buffer[:, conv_index, ...].fill_(0) # #17: zero the FULL (mtp_step + 1)-row SSM block, not just canonical row +0, so a future # first-step verify reading offset>0 after fresh init never hits a never-written row (NaN). - self.req_to_ssm_state.buffer[:, ssm_start : ssm_start + (self.mtp_step + 1), ...].fill_(0) + self.req_to_ssm_state.buffer[:, ssm_start : ssm_start + self.ssm_slots_per_req, ...].fill_(0) if self.req_to_mtp_state_index is not None: self.req_to_mtp_state_index[req.req_idx] = 0 return @@ -66,6 +104,8 @@ def create_small_page_cache_manager(self, size: int): def save_big_page_states(self, b_req_idx: torch.Tensor, req_indexes: List[int], buffer_indexes: List[int]): from lightllm.common.basemodel.triton_kernel.linear_att_copy import copy_linear_att_state_to_kv_buffer + if self.replay_cache is not None: + self.replay_cache.materialize(b_req_idx) buffer_indexes = torch.tensor(buffer_indexes, dtype=torch.int32, device="cpu").cuda(non_blocking=True) state_cache_manager = self.big_page_buffers copy_linear_att_state_to_kv_buffer( @@ -75,15 +115,21 @@ def save_big_page_states(self, b_req_idx: torch.Tensor, req_indexes: List[int], gpu_ssm_state=self.req_to_ssm_state.buffer, cpu_kv_conv_state=state_cache_manager.conv_state_cache.buffer, cpu_kv_ssm_state=state_cache_manager.ssm_state_cache.buffer, - mtp_step=self.mtp_step, + mtp_step=self.ssm_slots_per_req - 1, + conv_offsets=self.req_to_mtp_state_index if self.replay_cache is not None else None, ) return def save_state(self, req_idx: int, buffer_idx: int, state_cache_manager: LinearAttCacheManager): + if self.replay_cache is not None: + self.replay_cache.materialize(torch.tensor([req_idx], dtype=torch.int32, device="cuda")) # checkpoint 只保存标准 conv 窗口和请求的基准 SSM 状态，不包含 MTP 扩展运行态。 conv_cache_width = self.linear_config.get_conv_state_shape()[-1] gpu_conv_state = self.req_to_conv_state.buffer[:, req_idx, ..., :conv_cache_width] - gpu_ssm_state = self.req_to_ssm_state.buffer[:, req_idx * (self.mtp_step + 1), ...] + if self.replay_cache is not None and self.req_to_mtp_state_index is not None: + offsets = torch.arange(conv_cache_width, device="cuda") + self.req_to_mtp_state_index[req_idx] + gpu_conv_state = self.req_to_conv_state.buffer[:, req_idx].index_select(-1, offsets) + gpu_ssm_state = self.req_to_ssm_state.buffer[:, req_idx * self.ssm_slots_per_req, ...] dst_conv_state, dst_ssm_state = state_cache_manager.get_state_cache(buffer_idx=buffer_idx) dst_conv_state.copy_(gpu_conv_state, non_blocking=True) dst_ssm_state.copy_(gpu_ssm_state, non_blocking=True) @@ -109,10 +155,16 @@ def update_mtp_state(self, b_req_mtp_start_loc, b_req_idx, b_mtp_index, accepted verify_width=verify_width, ) + if self.replay_cache is not None: + reqs = b_req_idx[b_req_mtp_start_loc.long()] + self.replay_cache.commit(reqs, self.req_to_mtp_state_index) + def restore_state(self, req: "InferReq", state_cache_manager: LinearAttCacheManager, buffer_idx: int): + if self.replay_cache is not None: + self.replay_cache.reset(req.req_idx) conv_state, ssm_state = state_cache_manager.get_state_cache(buffer_idx=buffer_idx) conv_dest = req.req_idx - ssm_dest = req.req_idx * (self.mtp_step + 1) + ssm_dest = req.req_idx * self.ssm_slots_per_req conv_cache_width = conv_state.shape[-1] self.req_to_conv_state.buffer[:, conv_dest, ..., :conv_cache_width] = conv_state self.req_to_ssm_state.buffer[:, ssm_dest, ...] = ssm_state diff --git a/lightllm/server/api_cli.py b/lightllm/server/api_cli.py index 926f9f2030..27193ed6f3 100644 --- a/lightllm/server/api_cli.py +++ b/lightllm/server/api_cli.py @@ -940,6 +940,19 @@ def add_cli_args(parser: argparse.ArgumentParser) -> argparse.ArgumentParser: If radix cache hit rate is low under high load due to limited small-page capacity and LRU eviction, increasing linear_att_cache_size can improve hit rate at the cost of more memory.""", ) + parser.add_argument( + "--enable_replayssm", + action="store_true", + default=False, + help="Use ReplaySSM for GDN: deferred FP32 decode/MTP or compact BF16 MTP.", + ) + parser.add_argument( + "--replayssm_cache_len", + type=int, + default=16, + choices=[16, 32, 64], + help="Accepted-history capacity for ReplaySSM; must cover the MTP verify width.", + ) parser.add_argument( "--linear_att_ssm_data_type", type=str, diff --git a/lightllm/server/api_start.py b/lightllm/server/api_start.py index fea827d466..c8d7a5140a 100644 --- a/lightllm/server/api_start.py +++ b/lightllm/server/api_start.py @@ -279,6 +279,18 @@ def _launch_subprocesses(args: StartArgs): f"but got {args.batch_max_tokens}, {args.chunked_prefill_size}" ) + if getattr(args, "enable_replayssm", False): + from lightllm.utils.config_utils import get_model_type + + assert get_model_type(args.model_dir) in ( + "qwen3_next", + "qwen3_5", + "qwen3_5_moe", + "qwen3_5_text", + "qwen3_5_moe_text", + ), "ReplaySSM currently requires a GDN model" + assert args.replayssm_cache_len >= args.mtp_step + 1, "ReplaySSM capacity must cover the verify width" + # hybrid checkpoint 参数自动设置；保留现有 linear_att_* 启动参数名。 if args.linear_att_cache_size is None: # 小页池大小只对 hybrid 模型生效。 diff --git a/lightllm/server/core/objs/start_args_type.py b/lightllm/server/core/objs/start_args_type.py index 13ad6367a5..d78e38bcc2 100644 --- a/lightllm/server/core/objs/start_args_type.py +++ b/lightllm/server/core/objs/start_args_type.py @@ -242,6 +242,8 @@ class StartArgs: weight_version: str = "default" # hybrid attention model (Qwen3Next) + enable_replayssm: bool = False + replayssm_cache_len: int = 16 linear_att_hash_page_size: int = field(default=512) linear_att_page_block_num: int = field(default=10000000) disable_linear_att_small_page_cpu_cache: bool = field(default=False) diff --git a/test/benchmark/kernels/benchmark_replayssm.py b/test/benchmark/kernels/benchmark_replayssm.py new file mode 100644 index 0000000000..33357c986b --- /dev/null +++ b/test/benchmark/kernels/benchmark_replayssm.py @@ -0,0 +1,141 @@ +"""Run with PYTHONPATH=. exp -m 'ReplaySSM kernel comparison' python test/benchmark/kernels/benchmark_replayssm.py.""" +import argparse +import json + +import torch +import triton + +from lightllm.common.basemodel.triton_kernel.linear_att.replayssm import ReplaySSMCache +from lightllm.common.basemodel.triton_kernel.linear_att.mtp_fused_recurrent import ( + mtp_fused_recurrent_gated_delta_rule, +) +from lightllm.common.basemodel.triton_kernel.linear_att.fla.ops.fused_recurrent import ( + fused_recurrent_gated_delta_rule, +) + + +def main(): + parser = argparse.ArgumentParser() + parser.add_argument("--batches", nargs="+", type=int, default=[1, 16, 64, 256]) + parser.add_argument("--widths", nargs="+", type=int, default=[1, 3, 5]) + parser.add_argument("--compact", action="store_true") + parser.add_argument("--state-dtype", choices=["float32", "bfloat16"], default="float32") + parser.add_argument("--bv", type=int, default=64) + parser.add_argument("--warps", type=int, default=2) + parser.add_argument("--capacity", type=int, default=16) + parser.add_argument("--heads", type=int, default=16) + parser.add_argument("--value-heads", type=int) + parser.add_argument("--layers", type=int, default=1) + parser.add_argument("--accepted-tokens", type=int) + args = parser.parse_args() + torch.manual_seed(1) + h, hv, kd, vd = args.heads, args.heads * 2 if args.value_heads is None else args.value_heads, 128, 128 + assert h > 0 and hv > 0 and hv % h == 0 and args.layers > 0 + for width in args.widths: + accepted_tokens = width if args.accepted_tokens is None else args.accepted_tokens + assert 1 <= accepted_tokens <= width + for batch in args.batches: + q = torch.randn(1, batch * width, h, kd, dtype=torch.bfloat16, device="cuda") + k = torch.randn_like(q) + v = torch.randn(1, batch * width, hv, vd, dtype=torch.bfloat16, device="cuda") + a = torch.randn(batch * width, hv, device="cuda", dtype=torch.bfloat16) - 3 + b = torch.randn_like(a) + alog = torch.zeros(hv, device="cuda") + bias = torch.zeros_like(alog) + state = torch.zeros( + args.layers, batch + 1, hv, kd, vd, device="cuda", dtype=getattr(torch, args.state_dtype) + ) + if args.compact: + from lightllm.common.basemodel.triton_kernel.linear_att.replayssm_compact import CompactSSMCache + + replay = CompactSSMCache(state, width, torch.bfloat16) + else: + replay = ReplaySSMCache(state, args.capacity, width) + reqs = torch.arange(batch, device="cuda", dtype=torch.int32) + cu = torch.arange(batch + 1, device="cuda", dtype=torch.int32) * width + accepted = torch.full((batch + 1,), accepted_tokens - 1, device="cuda", dtype=torch.int32) + baseline_states = torch.zeros( + args.layers, (batch + 1) * width, hv, kd, vd, device="cuda", dtype=state.dtype + ) + indexes = torch.arange(batch * width, device="cuda", dtype=torch.int32).view(batch, width) + counts = torch.full((batch,), accepted_tokens, device="cuda", dtype=torch.int32) + + def run_replay(): + positions = replay.positions(reqs) + kwargs = {} if args.compact else {"run_config": {"BV": args.bv, "num_warps": args.warps}} + for layer in range(args.layers): + replay.forward(layer, q, k, v, a, b, alog, bias, reqs, positions, cu, **kwargs) + replay.commit(reqs, accepted) + + def run_baseline_layer(baseline_state): + if width > 1: + mtp_fused_recurrent_gated_delta_rule( + q, + k, + v, + baseline_state, + cu, + indexes, + indexes, + counts, + alog, + bias, + a, + b, + run_config={"BV": 8, "num_warps": 1, "num_stages": 1}, + ) + else: + fused_recurrent_gated_delta_rule( + q=q.view(batch, 1, h, kd), + k=k.view(batch, 1, h, kd), + v=v.view(batch, 1, hv, vd), + initial_state=baseline_state, + inplace_final_state=True, + ssm_state_indices=reqs, + use_qk_l2norm_in_kernel=True, + A_log=alog, + dt_bias=bias, + a_raw=a, + b_raw=b, + ) + + def run_baseline(): + for layer in range(args.layers): + run_baseline_layer(baseline_states[layer]) + + # Warm every ring phase before timing. Timed graphs advance real + # cursors too, including folds and metadata/acceptance launches. + for _ in range(2 * args.capacity): + run_replay() + run_baseline() + base_ms = triton.testing.do_bench_cudagraph(run_baseline, rep=300) + replay_ms = triton.testing.do_bench_cudagraph(run_replay, rep=300) + replay_bytes = sum( + x.numel() * x.element_size() for x in vars(replay).values() if isinstance(x, torch.Tensor) + ) + print( + json.dumps( + dict( + batch=batch, + width=width, + heads=hv, + layers=args.layers, + accepted_tokens=accepted_tokens, + compact=args.compact, + state_dtype=args.state_dtype, + capacity=args.capacity, + bv=8 if args.compact else args.bv, + warps=1 if args.compact else args.warps, + baseline_ms=base_ms, + replay_ms=replay_ms, + speedup=base_ms / replay_ms, + baseline_ssm_bytes=baseline_states.numel() * baseline_states.element_size(), + replay_bytes=replay_bytes, + ) + ), + flush=True, + ) + + +if __name__ == "__main__": + main() diff --git a/test/benchmark/service/compare_replayssm.py b/test/benchmark/service/compare_replayssm.py new file mode 100644 index 0000000000..feea74c503 --- /dev/null +++ b/test/benchmark/service/compare_replayssm.py @@ -0,0 +1,77 @@ +"""Small deterministic serving comparison; run through exp -m. + +This checks generation and prefix-cache behavior, not a task-accuracy benchmark. +Use identical model, prefill backend and sampling settings on both endpoints. +""" +import argparse +import json +import urllib.request +import concurrent.futures +import time +import pathlib + +parser = argparse.ArgumentParser() +parser.add_argument("--baseline-port", type=int, default=18932) +parser.add_argument("--replay-port", type=int, default=18934) +args = parser.parse_args() +prompts = [ + "The capital of France is", + "Write a Python function to add two numbers.", + "Count from one to twenty:", + "Explain why the sky is blue in one paragraph.", + "Here is some background. " + + ("The quick brown fox jumps over the lazy dog. " * 90) + + "\nSummarize this in a sentence:", +] + + +def run(arg): + port, prompt = arg + data = json.dumps( + {"inputs": prompt, "parameters": {"max_new_tokens": 128, "do_sample": False, "return_details": True}} + ).encode() + start = time.monotonic() + with urllib.request.urlopen( + urllib.request.Request( + f"http://127.0.0.1:{port}/generate", data=data, headers={"Content-Type": "application/json"} + ), + timeout=120, + ) as r: + res = json.load(r) + return { + "port": port, + "tokens": [t["id"] for t in res["tokens"]], + "text": res["generated_text"][0], + "prompt_tokens": res["prompt_tokens"], + "cache_len": res["tokens"][0].get("prompt_cache_len"), + "mtp_accepted": res["tokens"][-1].get("mtp_accepted_token_num"), + "elapsed": time.monotonic() - start, + } + + +results = [] +for round in range(2): + with concurrent.futures.ThreadPoolExecutor(max_workers=5) as pool: + baseline = list(pool.map(run, [(args.baseline_port, p) for p in prompts])) + replay = list(pool.map(run, [(args.replay_port, p) for p in prompts])) + for i, (a, b) in enumerate(zip(baseline, replay)): + results.append({"round": round, "prompt": i, "baseline": a, "replay": b}) + prefix = next( + (j for j, (x, y) in enumerate(zip(a["tokens"], b["tokens"])) if x != y), + min(len(a["tokens"]), len(b["tokens"])), + ) + print( + json.dumps( + { + "round": round, + "prompt": i, + "tokens_equal": a["tokens"] == b["tokens"], + "common_prefix": prefix, + "replay_cache": b["cache_len"], + "prompt_tokens": b["prompt_tokens"], + "mtp_accepted": b["mtp_accepted"], + } + ), + flush=True, + ) +pathlib.Path("/tmp/replay-comparison-results.json").write_text(json.dumps(results)) diff --git a/unit_tests/common/basemodel/triton_kernel/linear_att/test_replayssm.py b/unit_tests/common/basemodel/triton_kernel/linear_att/test_replayssm.py new file mode 100644 index 0000000000..41ab2afbbb --- /dev/null +++ b/unit_tests/common/basemodel/triton_kernel/linear_att/test_replayssm.py @@ -0,0 +1,320 @@ +import pytest +import torch + +from lightllm.common.basemodel.triton_kernel.linear_att.replayssm import ReplaySSMCache + +pytestmark = pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") + + +def reference(q, k, v, a, b, a_log, bias, state): + q, k, v = q.float(), k.float(), v.float() + q = q / (q.square().sum(-1, keepdim=True) + 1e-6).sqrt() * (q.shape[-1] ** -0.5) + k = k / (k.square().sum(-1, keepdim=True) + 1e-6).sqrt() + q = q.repeat_interleave(v.shape[-2] // q.shape[-2], -2) + k = k.repeat_interleave(v.shape[-2] // k.shape[-2], -2) + g = -a_log.exp() * torch.nn.functional.softplus(a.float() + bias) + state = state * g.exp()[..., None, None] + d = b.float().sigmoid()[..., None] * (v - torch.einsum("...kv,...k->...v", state, k)) + state = state + k[..., None] * d[..., None, :] + return torch.einsum("...kv,...k->...v", state, q), state + + +@pytest.mark.parametrize("capacity", [16, 32, 64]) +@pytest.mark.parametrize("width", [1, 3, 5, 16]) +@pytest.mark.parametrize("dims", [(32, 64), (128, 128)]) +def test_replay_acceptance_flush_and_graph(width, dims, capacity): + torch.manual_seed(21) + kdim, vdim = dims + h, hv, slots = 2, 4, 5 + state = torch.randn(2, slots, hv, kdim, vdim, device="cuda") * 0.01 + cache = ReplaySSMCache(state, capacity, width) + expected = state.clone() + reqs = torch.tensor([2, 0, slots - 1], device="cuda", dtype=torch.int32) + cu = torch.tensor([0, width, 2 * width - 1, 2 * width - 1], device="cuda", dtype=torch.int32) + if width == 1: + cu = torch.tensor([0, 1, 2, 2], device="cuda", dtype=torch.int32) + tokens = 3 * width + q = torch.randn(1, tokens, h, kdim, device="cuda", dtype=torch.bfloat16) + k = torch.randn_like(q) + v = torch.randn(1, tokens, hv, vdim, device="cuda", dtype=torch.bfloat16) + a = torch.randn(tokens, hv, device="cuda", dtype=torch.bfloat16) - 3 + b = torch.randn_like(a) + alog = torch.randn(hv, device="cuda") * 0.1 + bias = torch.randn(hv, device="cuda") * 0.1 + accepted = torch.zeros(slots, dtype=torch.int32, device="cuda") + + def step(): + pos = cache.positions(reqs) + out = [cache.forward(layer, q, k, v, a, b, alog, bias, reqs, pos, cu) for layer in range(2)] + cache.commit(reqs, accepted) + return out + + # Compile on disposable state, then capture with the HOLD slot only. + saved = state.clone() + step() + state.copy_(saved) + cache.cursors.zero_() + reqs.fill_(slots - 1) + torch.cuda.synchronize() + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + outputs = step() + reqs.copy_(torch.tensor([2, 0, slots - 1], device="cuda", dtype=torch.int32)) + for iteration in range(3 * capacity + 5): + counts = [iteration % width + 1, min((iteration + 1) % width + 1, max(width - 1, 1))] + accepted[2] = counts[0] - 1 + accepted[0] = counts[1] - 1 + graph.replay() + for layer in range(2): + for seq, req in enumerate([2, 0]): + start = seq * width + length = width if seq == 0 else max(width - 1, 1) + cur = expected[layer, req].clone() + for j in range(length): + t = start + j + out, cur = reference(q[0, t], k[0, t], v[0, t], a[t], b[t], alog, bias, cur) + torch.testing.assert_close(outputs[layer][0, t].float(), out, atol=0.006, rtol=0.02) + if j + 1 == counts[seq]: + expected[layer, req].copy_(cur) + if iteration == 2 * capacity + 1: + cache.materialize(reqs[:2]) + torch.testing.assert_close(state[:, :4], expected[:, :4], atol=2e-5, rtol=2e-4) + cache.materialize(reqs[:2]) + torch.testing.assert_close(state, expected, atol=2e-5, rtol=2e-4) + + +@pytest.mark.parametrize("dtype", [torch.float32, torch.bfloat16]) +@pytest.mark.parametrize("kind", ["gdn", "kda"]) +def test_compact_preserves_accepted_prefix(dtype, kind): + from lightllm.common.basemodel.triton_kernel.linear_att.replayssm_compact import CompactSSMCache + + torch.manual_seed(17) + hv, kd, vd, width = 4, 64, 64, 3 + state = torch.randn(2, 4, hv, kd, vd, device="cuda", dtype=dtype) * 0.01 + cache = CompactSSMCache(state, width, torch.bfloat16, kind) + reqs = torch.tensor([2, 0, 3], device="cuda", dtype=torch.int32) + cu = torch.tensor([0, 3, 5, 5], device="cuda", dtype=torch.int32) + accepted = torch.tensor([1, 0, 0, 0], device="cuda", dtype=torch.int32) + shape = (1, 5, hv, kd) + q = torch.randn(shape, device="cuda", dtype=torch.bfloat16) + k, v = torch.randn_like(q), torch.randn_like(q) + a = torch.randn((5, hv * kd if kind == "kda" else hv), device="cuda", dtype=torch.bfloat16) + b = torch.randn(5, hv, device="cuda", dtype=torch.bfloat16) + alog = torch.randn(hv, device="cuda") * 0.1 + bias = torch.randn(hv * kd if kind == "kda" else hv, device="cuda") * 0.1 + for _ in range(15): + before = state.clone() + outputs = [cache.forward(layer, q, k, v, a, b, alog, bias, reqs, None, cu) for layer in range(2)] + assert torch.equal(state, before), "verify must not mutate the checkpoint" + cache.commit(reqs, accepted) + for layer in range(2): + for req, start, length in [(2, 0, 3), (0, 3, 2)]: + cur = before[layer, req].float() + for j in range(length): + t = start + j + if kind == "gdn": + out, cur = reference(q[0, t], k[0, t], v[0, t], a[t], b[t], alog, bias, cur) + else: + qq, kk = q[0, t].float(), k[0, t].float() + qq *= torch.rsqrt(qq.square().sum(-1, keepdim=True) + 1e-6) * kd**-0.5 + kk *= torch.rsqrt(kk.square().sum(-1, keepdim=True) + 1e-6) + gate = -5 * torch.sigmoid(alog.exp()[:, None] * (a[t].float().view(hv, kd) + bias.view(hv, kd))) + cur *= gate.exp()[..., None] + d = (v[0, t].float() - torch.einsum("hkv,hk->hv", cur, kk)) * b[t].float().sigmoid()[:, None] + cur += kk[..., None] * d[:, None, :] + out = torch.einsum("hkv,hk->hv", cur, qq) + torch.testing.assert_close(outputs[layer][0, t].float(), out, atol=0.004, rtol=0.02) + cur = cur.to(dtype).float() + if j == int(accepted[req]): + torch.testing.assert_close( + state[layer, req].float(), cur, atol=0.008 if dtype == torch.bfloat16 else 2e-5, rtol=0.02 + ) + torch.testing.assert_close(state[:, 1], before[:, 1], rtol=0, atol=0) + torch.testing.assert_close(state[:, 3], before[:, 3], rtol=0, atol=0) + + +@pytest.mark.parametrize("dtype", [torch.float32, torch.bfloat16]) +@pytest.mark.parametrize("kind", ["gdn", "kda"]) +@pytest.mark.parametrize("width", [3, 4]) +def test_compact_matches_native_mtp(dtype, kind, width): + from lightllm.common.basemodel.triton_kernel.linear_att.replayssm_compact import CompactSSMCache + from lightllm.common.basemodel.triton_kernel.linear_att.mtp_fused_recurrent import ( + mtp_fused_recurrent_gated_delta_rule, + ) + + torch.manual_seed(77) + batch, h, hv, kd, vd = 3, (4 if kind == "kda" else 2), 4, 128, 128 + q = torch.randn(1, batch * width, h, kd, device="cuda", dtype=torch.bfloat16) + k = torch.randn_like(q) + v = torch.randn(1, batch * width, hv, vd, device="cuda", dtype=torch.bfloat16) + a = torch.randn(batch * width, hv * kd if kind == "kda" else hv, device="cuda", dtype=torch.bfloat16) + b = torch.randn(batch * width, hv, device="cuda", dtype=torch.bfloat16) + alog = torch.randn(hv, device="cuda") * 0.1 + bias = torch.randn(hv * kd if kind == "kda" else hv, device="cuda") * 0.1 + state = torch.randn(1, batch + 1, hv, kd, vd, device="cuda", dtype=dtype) * 0.01 + reference_state = state[0].repeat_interleave(width, 0) + reqs = torch.arange(batch, device="cuda", dtype=torch.int32) + cu = torch.arange(batch + 1, device="cuda", dtype=torch.int32) * width + idx = torch.arange(batch * width, device="cuda", dtype=torch.int32).view(batch, width) + counts = torch.ones(batch, device="cuda", dtype=torch.int32) + accepted = torch.tensor([0, width // 2, width - 1, 0], device="cuda", dtype=torch.int32) + cache = CompactSSMCache(state, width, torch.bfloat16, kind) + for _ in range(32): + q.normal_() + k.normal_() + v.normal_() + a.normal_().sub_(3) + b.normal_() + out = cache.forward(0, q, k, v, a, b, alog, bias, reqs, None, cu) + if kind == "kda": + native = pytest.importorskip("lightllm.common.basemodel.triton_kernel.linear_att.fla.ops.kda_decode") + ref, _ = native.fused_recurrent_kda( + q, + k, + v, + a.unsqueeze(0), + b.unsqueeze(0), + alog, + bias, + reference_state, + idx, + cu_seqlens=cu, + num_accepted_tokens=counts, + ) + else: + ref, _ = mtp_fused_recurrent_gated_delta_rule( + q, + k, + v, + reference_state, + cu, + idx, + idx, + counts, + alog, + bias, + a, + b, + run_config={"BV": 8, "num_warps": 1, "num_stages": 1}, + ) + cache.commit(reqs, accepted) + if kind == "gdn": + assert torch.equal(out, ref) + atol, rtol = 0, 0 + else: + # KDA's separate gate storage changes compiler contraction. Across + # repeated commits this can cross a BF16 rounding boundary. + torch.testing.assert_close(out, ref, atol=1e-5, rtol=1e-2) + atol, rtol = (1e-7, 2e-5) if dtype == torch.float32 else (1e-5, 1e-2) + torch.testing.assert_close( + state[0, :batch], reference_state[reqs * width + accepted[:batch]], rtol=rtol, atol=atol + ) + counts.copy_(accepted[:batch] + 1) + accepted[:batch].add_(1).remainder_(width) + + +@pytest.mark.parametrize("save_big", [False, True]) +def test_checkpoint_restores_pending_history_and_accepted_conv(save_big, monkeypatch): + from types import SimpleNamespace + from lightllm.common.req_manager.linear_att import ReqManagerForMamba + from lightllm.common.state_cache_manager import LinearAttCacheManager, LayerCache + from lightllm.common.kv_cache_mem_manager.qwen3next_mem_manager import Qwen3NextLinearAttPageHelper + + # Real save/restore and PD helpers, with just the unrelated allocator omitted. + manager = object.__new__(ReqManagerForMamba) + config = SimpleNamespace( + linear_layer_num=2, + conv_state_dtype=torch.bfloat16, + ssm_state_dtype=torch.float32, + get_conv_state_shape=lambda: (192, 3), + get_ssm_state_shape=lambda: (2, 32, 64), + global_linear_k_heads=1, + global_linear_v_heads=2, + head_linear_k_dim=32, + head_linear_v_dim=64, + num_linear_k_heads=1, + num_linear_v_heads=2, + tp_world_size=1, + conv_kernel_size=4, + ) + manager.linear_config = config + manager.mtp_step = 2 + manager.ssm_slots_per_req = 1 + manager.req_to_mtp_state_index = torch.zeros(4, device="cuda", dtype=torch.int32) + manager.req_to_mtp_state_index[1] = 1 + manager.req_to_conv_state = LayerCache(4, torch.bfloat16, (192, 5), 2, "cuda") + manager.req_to_conv_state.buffer.copy_(torch.randn_like(manager.req_to_conv_state.buffer)) + manager.req_to_ssm_state = LayerCache(4, torch.float32, (2, 32, 64), 2, "cuda") + manager.req_to_ssm_state.buffer.zero_() + manager.replay_cache = ReplaySSMCache(manager.req_to_ssm_state.buffer, 16, 3) + cache = manager.replay_cache + cache.raw_keys.normal_() + cache.raw_values.normal_() + cache.gates.uniform_(-0.1, -0.01) + cache.betas.uniform_(0.1, 0.9) + cache.cursors[1] = 2 + expected = manager.req_to_ssm_state.buffer[:, 1].clone() + for i in range(2): + key = cache.raw_keys[:, 1, :, i].float() + key /= (key.square().sum(-1, keepdim=True) + 1e-6).sqrt() + expected *= cache.gates[:, 1, :, i].exp()[..., None, None] + delta = cache.betas[:, 1, :, i, None] * ( + cache.raw_values[:, 1, :, i].float() - torch.einsum("lhkv,lhk->lhv", expected, key) + ) + expected += key[..., None] * delta[..., None, :] + conv_expected = manager.req_to_conv_state.buffer[:, 1, :, 1:4].clone() + cpu_cache = LinearAttCacheManager(2, config) + manager.mem_manager = SimpleNamespace(big_page_buffers=cpu_cache) + reqs = torch.tensor([1], device="cuda", dtype=torch.int32) + if save_big: + manager.save_big_page_states(reqs, [1], [0]) + else: + manager.save_state(1, 0, cpu_cache) + torch.cuda.synchronize() + conv_saved, state_saved = cpu_cache.get_state_cache(0) + torch.testing.assert_close(state_saved, expected.cpu(), rtol=2e-5, atol=2e-5) + torch.testing.assert_close(conv_saved, conv_expected.cpu(), rtol=0, atol=0) + cache.cursors[2] = 13 # stale history from a previous owner must disappear + manager.restore_state(SimpleNamespace(req_idx=2), cpu_cache, 0) + assert cache.cursors[2].item() == 0 + assert manager.req_to_mtp_state_index[2].item() == 0 + torch.testing.assert_close(manager.req_to_ssm_state.buffer[:, 2], expected, rtol=2e-5, atol=2e-5) + torch.testing.assert_close(manager.req_to_conv_state.buffer[:, 2, :, :3], conv_expected, rtol=0, atol=0) + + import lightllm.common.kv_cache_mem_manager.qwen3next_mem_manager as memory_module + + monkeypatch.setattr(memory_module, "get_env_start_args", lambda: SimpleNamespace(mtp_step=2)) + mem = SimpleNamespace( + linear_config=config, + req_to_conv_state=manager.req_to_conv_state, + req_to_ssm_state=manager.req_to_ssm_state, + replay_cache=cache, + ssm_slots_per_req=1, + req_to_mtp_state_index=manager.req_to_mtp_state_index, + ) + helper = Qwen3NextLinearAttPageHelper(mem) + # A second pending suffix exercises PD export, independently of CPU save. + cache.raw_keys[:, 1, :, 0].normal_() + cache.raw_values[:, 1, :, 0].normal_() + cache.gates[:, 1, :, 0].uniform_(-0.1, -0.01) + cache.betas[:, 1, :, 0].uniform_(0.1, 0.9) + cache.cursors[1] = 1 + key = cache.raw_keys[:, 1, :, 0].float() + key /= (key.square().sum(-1, keepdim=True) + 1e-6).sqrt() + expected *= cache.gates[:, 1, :, 0].exp()[..., None, None] + delta = cache.betas[:, 1, :, 0, None] * ( + cache.raw_values[:, 1, :, 0].float() - torch.einsum("lhkv,lhk->lhv", expected, key) + ) + expected += key[..., None] * delta[..., None, :] + conv_page = torch.empty(helper.conv_shape, device="cuda", dtype=torch.bfloat16) + ssm_page = torch.empty(helper.ssm_shape, device="cuda") + helper._write_one_rank(mem, 0, 1, conv_page, ssm_page) + cache.cursors[2] = 7 + manager.req_to_mtp_state_index[2] = 2 + helper._read_one_rank(mem, 0, 2, conv_page, ssm_page) + torch.testing.assert_close(manager.req_to_ssm_state.buffer[:, 2], expected, rtol=2e-5, atol=2e-5) + torch.testing.assert_close(manager.req_to_conv_state.buffer[:, 2, :, :3], conv_expected, rtol=0, atol=0) + assert cache.cursors[2].item() == manager.req_to_mtp_state_index[2].item() == 0 + manager.init_hybrid_attention_state(SimpleNamespace(req_idx=2)) + assert torch.count_nonzero(manager.req_to_ssm_state.buffer[:, 2]).item() == 0 + assert torch.count_nonzero(manager.req_to_conv_state.buffer[:, 2]).item() == 0