From 762c235848121e12d13cf6a7562e31a0bed711b5 Mon Sep 17 00:00:00 2001 From: sufubao Date: Sun, 20 Sep 2026 17:06:24 +0800 Subject: [PATCH 01/40] 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 From 4213a743b4365438a77dcc02d0d81aa1bc43635f Mon Sep 17 00:00:00 2001 From: sufubao Date: Tue, 6 Oct 2026 11:57:19 +0800 Subject: [PATCH 02/40] style: use the pinned Black formatting --- .../common/basemodel/triton_kernel/linear_att/replayssm.py | 4 ++-- .../basemodel/triton_kernel/linear_att/replayssm_compact.py | 4 ++-- lightllm/common/kv_cache_mem_manager/mem_manager.py | 2 +- .../basemodel/triton_kernel/linear_att/test_replayssm.py | 2 +- .../triton_kernel/linear_att/test_replayssm_execution.py | 4 ++-- 5 files changed, 8 insertions(+), 8 deletions(-) diff --git a/lightllm/common/basemodel/triton_kernel/linear_att/replayssm.py b/lightllm/common/basemodel/triton_kernel/linear_att/replayssm.py index 621522dfa7..43ac6c0716 100644 --- a/lightllm/common/basemodel/triton_kernel/linear_att/replayssm.py +++ b/lightllm/common/basemodel/triton_kernel/linear_att/replayssm.py @@ -240,7 +240,7 @@ def _replay( ).to(tl.float32) qk = q_rows + k_rows qk *= tl.rsqrt(tl.sum(qk * qk, axis=1) + 1.0e-6)[:, None] - qk *= tl.where(tt < BW, K**-0.5, 1.0)[:, None] + qk *= tl.where(tt < BW, K ** -0.5, 1.0)[:, None] projections = tl.dot(qk.to(projection_dtype), state.to(projection_dtype), input_precision=projection_precision) # Project the whole verify window against both accepted and new keys. @@ -350,7 +350,7 @@ def _replay( 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) + 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) diff --git a/lightllm/common/basemodel/triton_kernel/linear_att/replayssm_compact.py b/lightllm/common/basemodel/triton_kernel/linear_att/replayssm_compact.py index 996a8cd73f..55b3b65bb0 100644 --- a/lightllm/common/basemodel/triton_kernel/linear_att/replayssm_compact.py +++ b/lightllm/common/basemodel/triton_kernel/linear_att/replayssm_compact.py @@ -111,9 +111,9 @@ def _compact( 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)) + 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) + 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: diff --git a/lightllm/common/kv_cache_mem_manager/mem_manager.py b/lightllm/common/kv_cache_mem_manager/mem_manager.py index 9eecf369e0..eb8aee4e96 100755 --- a/lightllm/common/kv_cache_mem_manager/mem_manager.py +++ b/lightllm/common/kv_cache_mem_manager/mem_manager.py @@ -97,7 +97,7 @@ def profile_size(self, mem_fraction): available_memory = get_available_gpu_memory(world_size) - get_total_gpu_memory() * (1 - mem_fraction) cell_size = self.get_cell_size() pd_kv_move_buffer_size = self.get_pd_kv_move_buffer_size() - available_memory_bytes = available_memory * 1024**3 - pd_kv_move_buffer_size + available_memory_bytes = available_memory * 1024 ** 3 - pd_kv_move_buffer_size self.size = int(available_memory_bytes / cell_size) if world_size > 1: tensor = torch.tensor(self.size, dtype=torch.int64, device=f"cuda:{get_current_device_id()}") 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 index 9764a7ff2f..1e51b6bb0a 100644 --- a/unit_tests/common/basemodel/triton_kernel/linear_att/test_replayssm.py +++ b/unit_tests/common/basemodel/triton_kernel/linear_att/test_replayssm.py @@ -394,7 +394,7 @@ def test_compact_preserves_accepted_prefix(dtype, kind, run_config=None): 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 + 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] diff --git a/unit_tests/common/basemodel/triton_kernel/linear_att/test_replayssm_execution.py b/unit_tests/common/basemodel/triton_kernel/linear_att/test_replayssm_execution.py index eba08c7a3f..dc2ddce177 100644 --- a/unit_tests/common/basemodel/triton_kernel/linear_att/test_replayssm_execution.py +++ b/unit_tests/common/basemodel/triton_kernel/linear_att/test_replayssm_execution.py @@ -209,12 +209,12 @@ def warm(cache): def test_cross_layer_fold_with_state_offset_above_int32(): - if torch.cuda.mem_get_info()[0] < 6 * 1024**3: + if torch.cuda.mem_get_info()[0] < 6 * 1024 ** 3: pytest.skip("large-stride regression requires 6 GiB of free device memory") torch.manual_seed(791) shape = (48, 3, 12, 128, 128) strides = (50_528_256, 196608, 16384, 128, 1) - assert strides[0] < 2**31 < 47 * strides[0] + assert strides[0] < 2 ** 31 < 47 * strides[0] dense = torch.randn(shape, device="cuda", dtype=torch.bfloat16) * 0.01 sparse = torch.empty_strided(shape, strides, device="cuda", dtype=dense.dtype) sparse.copy_(dense) From 2727aaef6ae827d9035932e902679bb93f3a6ce0 Mon Sep 17 00:00:00 2001 From: sufubao Date: Tue, 6 Oct 2026 12:03:07 +0800 Subject: [PATCH 03/40] fix: skip hybrid checkpoints when prefix caching is disabled --- .../server/router/model_infer/infer_batch.py | 2 +- .../basemodel/attention/linear/test_gdn.py | 23 +++++++++++++++++++ 2 files changed, 24 insertions(+), 1 deletion(-) diff --git a/lightllm/server/router/model_infer/infer_batch.py b/lightllm/server/router/model_infer/infer_batch.py index 99f58ad67b..773f446980 100644 --- a/lightllm/server/router/model_infer/infer_batch.py +++ b/lightllm/server/router/model_infer/infer_batch.py @@ -402,7 +402,7 @@ def get_can_alloc_token_num(self): def save_hybrid_state_to_cache(self, b_req_idx: torch.Tensor, reqs: List["InferReq"]): """Snapshot request-level attention state at big/small-page boundaries.""" - if not self.is_hybrid_att_model: + if not self.is_hybrid_att_model or self.radix_cache is None: return # Request-state snapshot at a big-page boundary. diff --git a/unit_tests/common/basemodel/attention/linear/test_gdn.py b/unit_tests/common/basemodel/attention/linear/test_gdn.py index 8aa48b3c8a..ab85bbea87 100644 --- a/unit_tests/common/basemodel/attention/linear/test_gdn.py +++ b/unit_tests/common/basemodel/attention/linear/test_gdn.py @@ -11,6 +11,7 @@ from lightllm.common.basemodel.attention.linear.flashqla import FlashQlaLinearAttBackend from lightllm.common.basemodel.attention.linear.triton import TritonLinearAttBackend from lightllm.common.req_manager import ReqManager, ReqManagerForMamba +from lightllm.server.router.model_infer.infer_batch import InferenceContext from lightllm.common.state_cache_manager import LinearAttCacheConfig from lightllm.server.api_cli import make_argument_parser import lightllm.utils.backend_validator as backend_validator @@ -408,3 +409,25 @@ def test_gdn_prefill_backend_tries_candidates_in_order(monkeypatch, auto_linear_ assert backend_class is flashqla3_backend assert validate_calls == ["flashqla2", "flashqla3"] + + +@pytest.mark.parametrize("input_len", [4096, 8192]) +def test_hybrid_checkpoint_is_noop_when_prompt_cache_disabled(input_len): + context = InferenceContext() + context.is_hybrid_att_model = True + context.radix_cache = None + context.args = SimpleNamespace( + linear_att_hash_page_size=1024, linear_att_page_block_num=8, disable_chunked_prefill=False + ) + req = SimpleNamespace( + req_idx=0, + get_chuncked_input_token_len=lambda: input_len, + hybrid_cache_len=input_len, + hybrid_len_to_big_page_id={}, + tail_small_page_buffer_id=None, + ) + + context.save_hybrid_state_to_cache(torch.tensor([0]), [req]) + + assert req.hybrid_len_to_big_page_id == {} + assert req.tail_small_page_buffer_id is None From b63c50b1505ee9ab405cdf8fc07b84e847f3acde Mon Sep 17 00:00:00 2001 From: sufubao Date: Tue, 6 Oct 2026 12:30:14 +0800 Subject: [PATCH 04/40] refactor: simplify GDN SSM state handling --- .../common/basemodel/attention/linear/gdn.py | 6 +- .../triton_kernel/linear_att/replayssm.py | 2 +- .../linear_att/replayssm_compact.py | 63 ++------- .../triton_kernel/linear_att/ssm_autotune.py | 10 +- lightllm/common/req_manager/linear_att.py | 8 +- .../basemodel/attention/linear/test_gdn.py | 2 +- .../linear_att/test_acceptance_fusion.py | 4 +- .../linear_att/test_replayssm.py | 121 ++++++------------ .../linear_att/test_replayssm_execution.py | 24 ++-- .../linear_att/test_replayssm_grouped_keys.py | 12 +- .../linear_att/test_ssm_autotune.py | 18 +-- 11 files changed, 89 insertions(+), 181 deletions(-) diff --git a/lightllm/common/basemodel/attention/linear/gdn.py b/lightllm/common/basemodel/attention/linear/gdn.py index bd639d0e9f..89981ad934 100644 --- a/lightllm/common/basemodel/attention/linear/gdn.py +++ b/lightllm/common/basemodel/attention/linear/gdn.py @@ -209,8 +209,6 @@ class LinearAttDecodeAttState(BaseDecodeAttState): b1_mtp_cu_q_seq_len: torch.Tensor = None b_num_accepted_tokens: torch.Tensor = None - b_ssm_history_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: @@ -221,7 +219,7 @@ def init_state(self): self._init_fixed_mtp_decode_state(draft_step) ssm_updates = self.backend.req_manager.ssm_update_cache if ssm_updates is not None: - self.b_ssm_history_positions = ssm_updates.prepare_decode(self.b_conv_buffer_idx, self.b1_mtp_cu_q_seq_len) + ssm_updates.prepare_decode(self.b_conv_buffer_idx, self.b1_mtp_cu_q_seq_len) if draft_step == 0: ssm_updates.accept_updates(self.b_conv_buffer_idx) @@ -370,7 +368,6 @@ def _gdn_decode_kernel( layer_weight.linear_A_log.weight, layer_weight.linear_dt_bias.weight, self.b_conv_buffer_idx, - self.b_ssm_history_positions, ), z, ) @@ -435,7 +432,6 @@ def _gdn_mtp_kernel( layer_weight.linear_A_log.weight, layer_weight.linear_dt_bias.weight, self.b_conv_buffer_idx, - self.b_ssm_history_positions, cu_seqlens_q, ) assert self.b_ssm_buffer_idx.dim() == 2, "SSM buffer idx must be 2D [N, S+1]" diff --git a/lightllm/common/basemodel/triton_kernel/linear_att/replayssm.py b/lightllm/common/basemodel/triton_kernel/linear_att/replayssm.py index 43ac6c0716..74ce36492e 100644 --- a/lightllm/common/basemodel/triton_kernel/linear_att/replayssm.py +++ b/lightllm/common/basemodel/triton_kernel/linear_att/replayssm.py @@ -595,7 +595,7 @@ def snapshot_accepted_state(self, req_idx): self._materialize_accepted_state(reqs, output, snapshot=True) return output - def forward(self, layer, q, k, v, a, b, a_log, bias, reqs, positions, cu_seqlens=None, run_config=None): + def forward(self, layer, q, k, v, a, b, a_log, bias, reqs, cu_seqlens=None, run_config=None): if layer == 0 and run_config is None: configure_cache(self, "replay", q, k, v, a, b, a_log, bias, cu_seqlens) hv, kd, vd = self.state.shape[-3:] diff --git a/lightllm/common/basemodel/triton_kernel/linear_att/replayssm_compact.py b/lightllm/common/basemodel/triton_kernel/linear_att/replayssm_compact.py index 55b3b65bb0..81e19c6d91 100644 --- a/lightllm/common/basemodel/triton_kernel/linear_att/replayssm_compact.py +++ b/lightllm/common/basemodel/triton_kernel/linear_att/replayssm_compact.py @@ -1,9 +1,8 @@ # SPDX-License-Identifier: Apache-2.0 -"""Compact speculative state for GDN/BF16 and KDA. +"""Compact speculative state for GDN. 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 @@ -47,8 +46,6 @@ def _compact( 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) @@ -75,45 +72,30 @@ def _compact( if COMMIT: k = tl.load(Keys + key_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) + decay = 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) + 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: if hv % (HV // KH) == 0: tl.store(Keys + key_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 + tl.store(Decays + record, decay) + 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) + 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: @@ -121,25 +103,12 @@ def _compact( class CompactSSMCache: - def __init__( - self, - state, - verify_width, - activation_dtype, - kind="gdn", - lower_bound=-5.0, - *, - num_key_heads=None, - run_config=None - ): - assert kind in ("gdn", "kda") + def __init__(self, state, verify_width, activation_dtype, *, num_key_heads=None, run_config=None): self.state = state self.verify_width = verify_width - self.kind = kind - self.lower_bound = lower_bound # Verify and accepted-state reconstruction must share the same layout. self._config_is_fixed = run_config is not None - self.run_config = dict(run_config or {"BV": 32 if kind == "kda" else 8, "num_warps": 4 if kind == "kda" else 1}) + self.run_config = dict(run_config or {"BV": 8, "num_warps": 1}) assert self.run_config["BV"] in (8, 16, 32, 64, 128) assert self.run_config["num_warps"] in (1, 2, 4, 8) layers, slots, hv, k, v = state.shape @@ -151,7 +120,7 @@ def __init__( (layers, slots, self.num_key_heads, verify_width, 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.decays = torch.empty(shape, device=state.device, dtype=torch.float32) self.betas = torch.empty(shape, device=state.device, dtype=torch.float32) def clear_history(self, req): @@ -169,10 +138,10 @@ def merge_accepted_updates(self, reqs): def snapshot_accepted_state(self, req_idx): return self.state[:, req_idx] - def forward(self, layer, q, k, v, a, b, a_log, bias, reqs, positions, cu_seqlens=None): + def forward(self, layer, q, k, v, a, b, a_log, bias, reqs, cu_seqlens=None): assert cu_seqlens is not None, "compact replay is only used for speculative verify" if layer == 0: - configure_cache(self, self.kind, q, k, v, a, b, a_log, bias, cu_seqlens) + configure_cache(self, "gdn", q, k, v, a, b, a_log, bias, cu_seqlens) _, slots, hv, kd, vd = self.state.shape assert self.num_key_heads in (q.shape[-2], hv) out = torch.empty_like(v) @@ -211,8 +180,6 @@ def forward(self, layer, q, k, v, a, b, a_log, bias, reqs, positions, cu_seqlens triton.next_power_of_2(kd), bv, False, - self.kind == "kda", - self.lower_bound, num_warps=self.run_config["num_warps"], num_stages=self.run_config.get("num_stages", 3), ) @@ -255,8 +222,6 @@ def accept_updates(self, reqs, accepted): triton.next_power_of_2(kd), bv, True, - self.kind == "kda", - self.lower_bound, num_warps=self.run_config["num_warps"], num_stages=self.run_config.get("num_stages", 3), ) diff --git a/lightllm/common/basemodel/triton_kernel/linear_att/ssm_autotune.py b/lightllm/common/basemodel/triton_kernel/linear_att/ssm_autotune.py index a155d5e022..6bb435cb77 100644 --- a/lightllm/common/basemodel/triton_kernel/linear_att/ssm_autotune.py +++ b/lightllm/common/basemodel/triton_kernel/linear_att/ssm_autotune.py @@ -28,7 +28,7 @@ def static_key(cache, mode, q, k, v, a, b, cu_seqlens): "dtype": str((q.dtype, cache.state.dtype)), "strides": str((q.stride(axis), k.stride(axis), v.stride(axis), a.stride(0), b.stride(0))), "history": cache.capacity if mode == "replay" else 0, - "gate": cache.lower_bound if mode == "kda" else 0, + "gate": 0, "varlen": cu_seqlens is not None, } if mode == "replay": @@ -61,9 +61,7 @@ def rebuild_inputs(cache, mode, q, k, v, a, b, a_log, bias, cu_seqlens=None, wor projection_mode=cache.projection_mode, ) else: - scratch = type(cache)( - state, cache.verify_width, q.dtype, mode, cache.lower_bound, num_key_heads=cache.num_key_heads - ) + scratch = type(cache)(state, cache.verify_width, q.dtype, num_key_heads=cache.num_key_heads) # Benchmark callbacks must not recursively select a configuration. scratch._config_is_fixed = True inputs = [ @@ -117,9 +115,9 @@ def select_config(cache, mode, q, k, v, a, b, a_log, bias, cu_seqlens=None, work else: rounds = 2 # Partial acceptance and full acceptance, both across all layers. for step in range(rounds): - positions = cache.prepare_decode(reqs, cu_seqlens) + cache.prepare_decode(reqs, cu_seqlens) for layer in range(cache.state.shape[0]): - cache.forward(layer, q, k, v, a, b, a_log, bias, reqs, positions, cu_seqlens) + cache.forward(layer, q, k, v, a, b, a_log, bias, reqs, cu_seqlens) cache.accept_updates(reqs, partial if mode != "replay" and step == 0 else accepted) if mode == "replay" and step == rounds - 2: cache._materialize_accepted_state(reqs[:1], snapshot, snapshot=True) diff --git a/lightllm/common/req_manager/linear_att.py b/lightllm/common/req_manager/linear_att.py index c46315ad1a..e8e1ab5170 100644 --- a/lightllm/common/req_manager/linear_att.py +++ b/lightllm/common/req_manager/linear_att.py @@ -19,7 +19,6 @@ def __init__( max_sequence_length, mem_manager, linear_config: LinearAttCacheConfig, - recurrent_kind="gdn", ): super().__init__(max_request_num, max_sequence_length, mem_manager) args = get_env_start_args() @@ -59,7 +58,6 @@ def __init__( if args.ssm_state_mode == "replay": from lightllm.common.basemodel.triton_kernel.linear_att.replayssm import ReplaySSMCache - assert recurrent_kind == "gdn", "ssm_state_mode=replay requires GDN" self.ssm_update_cache = ReplaySSMCache( self.req_to_ssm_state.buffer, args.replayssm_cache_len, @@ -75,7 +73,6 @@ def __init__( self.req_to_ssm_state.buffer, self.mtp_step + 1, linear_config.conv_state_dtype, - kind=recurrent_kind, num_key_heads=linear_config.num_linear_k_heads, ) return @@ -100,10 +97,7 @@ def save_big_page_states(self, b_req_idx: torch.Tensor, req_indexes: List[int], from lightllm.common.basemodel.triton_kernel.linear_att_copy import copy_linear_att_state_to_kv_buffer if self.ssm_update_cache is not None: - for req_idx, buffer_idx in zip(req_indexes, buffer_indexes): - if buffer_idx != -1: - self.save_state(req_idx, buffer_idx, self.big_page_buffers) - return + return super().save_big_page_states(b_req_idx, req_indexes, buffer_indexes) 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( diff --git a/unit_tests/common/basemodel/attention/linear/test_gdn.py b/unit_tests/common/basemodel/attention/linear/test_gdn.py index ab85bbea87..de4cfe5040 100644 --- a/unit_tests/common/basemodel/attention/linear/test_gdn.py +++ b/unit_tests/common/basemodel/attention/linear/test_gdn.py @@ -183,7 +183,7 @@ def prepare_decode(actual_reqs, actual_cu): accept_updates=lambda *args: pytest.fail("speculative updates must wait for acceptance"), ) state.init_state() - assert state.b_ssm_history_positions is positions + assert not hasattr(state, "b_ssm_history_positions") assert state.b_ssm_buffer_idx is reqs diff --git a/unit_tests/common/basemodel/triton_kernel/linear_att/test_acceptance_fusion.py b/unit_tests/common/basemodel/triton_kernel/linear_att/test_acceptance_fusion.py index a89f04e710..48de74ec06 100644 --- a/unit_tests/common/basemodel/triton_kernel/linear_att/test_acceptance_fusion.py +++ b/unit_tests/common/basemodel/triton_kernel/linear_att/test_acceptance_fusion.py @@ -136,8 +136,8 @@ def test_fused_acceptance_preserves_folds_snapshots_and_slot_reuse(dtype, projec def cycle(i): cache = caches[i] - positions = cache.prepare_decode(reqs, cu) - outputs = [cache.forward(layer, q, k, v, a, b, alog, alog, reqs, positions, cu) for layer in range(2)] + cache.prepare_decode(reqs, cu) + outputs = [cache.forward(layer, q, k, v, a, b, alog, alog, reqs, cu) for layer in range(2)] (old if i == 0 else new)(managers[i], args) return outputs 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 index 1e51b6bb0a..a224ed0400 100644 --- a/unit_tests/common/basemodel/triton_kernel/linear_att/test_replayssm.py +++ b/unit_tests/common/basemodel/triton_kernel/linear_att/test_replayssm.py @@ -70,8 +70,8 @@ def test_replay_acceptance_flush_and_graph( accepted = torch.zeros(slots, dtype=torch.int32, device="cuda") def step(): - pos = cache.prepare_decode(reqs, cu) - out = [cache.forward(layer, q, k, v, a, b, alog, bias, reqs, pos, cu) for layer in range(2)] + cache.prepare_decode(reqs, cu) + out = [cache.forward(layer, q, k, v, a, b, alog, bias, reqs, cu) for layer in range(2)] cache.accept_updates(reqs, accepted) return out @@ -217,8 +217,8 @@ def test_precomputed_state_preserves_projection_precision(warps, decay, state_dt accepted[1] = step % width outputs = [] for cache in caches: - positions = cache.prepare_decode(reqs, cu) - outputs.append(cache.forward(0, q, k, v, a, b, alog, alog, reqs, positions, cu)) + cache.prepare_decode(reqs, cu) + outputs.append(cache.forward(0, q, k, v, a, b, alog, alog, reqs, cu)) cache.accept_updates(reqs, accepted) if dtype == torch.float32: torch.testing.assert_close(outputs[0], outputs[1], atol=2e-6, rtol=2e-5) @@ -254,9 +254,9 @@ def test_bf16_decode_fold_uses_rounded_checkpoint(): bias = torch.zeros_like(alog) outputs = [] for cache in caches: - positions = cache.prepare_decode(reqs) + cache.prepare_decode(reqs) cache.accept_updates(reqs) - outputs.append(cache.forward(0, q, k, v, a, b, alog, bias, reqs, positions)) + outputs.append(cache.forward(0, q, k, v, a, b, alog, bias, reqs)) torch.testing.assert_close(outputs[0], outputs[1], atol=1e-7, rtol=1e-6) torch.testing.assert_close(caches[0].state, caches[1].state, atol=0, rtol=0) @@ -282,12 +282,12 @@ def test_dynamic_verify_lengths_fill_history_before_folding(dtype, projection_mo b = torch.randn_like(a) alog = torch.zeros(4, device="cuda") - def forward(positions): - return [cache.forward(layer, q, k, v, a, b, alog, alog, reqs, positions, cu) for layer in range(2)] + def forward(): + return [cache.forward(layer, q, k, v, a, b, alog, alog, reqs, cu) for layer in range(2)] def step(): positions = cache.prepare_decode(reqs, cu) - out = forward(positions) + out = forward() cache.accept_updates(reqs, accepted) return positions, out @@ -299,7 +299,7 @@ def step(): if capture_metadata: positions, outputs = step() else: - outputs = forward(positions) + outputs = forward() lengths = [(3, 4), (2, 1), (3, 2), (1, 4), (4, 3), (0, 2), (2, 0), (4, 4)] for iteration in range(24): ids = [0, 1] if iteration < 3 or iteration % 2 == 0 else [1, 0] @@ -362,27 +362,26 @@ def test_snapshot_without_history_copies_active_state(): @pytest.mark.parametrize("dtype", [torch.float32, torch.bfloat16]) -@pytest.mark.parametrize("kind", ["gdn", "kda"]) -def test_compact_preserves_accepted_prefix(dtype, kind, run_config=None): +def test_compact_preserves_accepted_prefix(dtype, run_config=None): 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, num_key_heads=hv, run_config=run_config) + cache = CompactSSMCache(state, width, torch.bfloat16, num_key_heads=hv, run_config=run_config) 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) + a = torch.randn((5, 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 + bias = torch.randn(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)] + outputs = [cache.forward(layer, q, k, v, a, b, alog, bias, reqs, cu) for layer in range(2)] assert torch.equal(state, before), "verify must not mutate the checkpoint" cache.accept_updates(reqs, accepted) for layer in range(2): @@ -390,17 +389,7 @@ def test_compact_preserves_accepted_prefix(dtype, kind, run_config=None): 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) + 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.004, rtol=0.02) cur = cur.to(dtype).float() if j == int(accepted[req]): @@ -412,23 +401,22 @@ def test_compact_preserves_accepted_prefix(dtype, kind, run_config=None): @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, run_config=None): +def test_compact_matches_native_mtp(dtype, width, run_config=None): 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 + batch, h, hv, kd, vd = 3, 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) + a = torch.randn(batch * width, 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 + bias = torch.randn(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) @@ -436,69 +424,42 @@ def test_compact_matches_native_mtp(dtype, kind, width, run_config=None): 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, num_key_heads=h, run_config=run_config) + cache = CompactSSMCache(state, width, torch.bfloat16, num_key_heads=h, run_config=run_config) 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={"num_stages": 1, **(run_config or {"BV": 8, "num_warps": 1})}, - ) - cache.accept_updates(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 + out = cache.forward(0, q, k, v, a, b, alog, bias, reqs, cu) + ref, _ = mtp_fused_recurrent_gated_delta_rule( + q, + k, + v, + reference_state, + cu, + idx, + idx, + counts, + alog, + bias, + a, + b, + run_config={"num_stages": 1, **(run_config or {"BV": 8, "num_warps": 1})}, ) + cache.accept_updates(reqs, accepted) + assert torch.equal(out, ref) + torch.testing.assert_close(state[0, :batch], reference_state[reqs * width + accepted[:batch]], rtol=0, atol=0) counts.copy_(accepted[:batch] + 1) accepted[:batch].add_(1).remainder_(width) @pytest.mark.parametrize("dtype", [torch.float32, torch.bfloat16]) @pytest.mark.parametrize("run_config", get_configs(), ids=lambda c: f"bv{c['BV']}-w{c['num_warps']}-s{c['num_stages']}") -@pytest.mark.parametrize("mode", ["gdn", "kda", "replay"]) +@pytest.mark.parametrize("mode", ["gdn", "replay"]) def test_autotune_candidates_match_recurrence(dtype, run_config, mode): if mode == "gdn": - test_compact_matches_native_mtp(dtype, "gdn", 4, run_config) - elif mode == "kda": - test_compact_preserves_accepted_prefix(dtype, "kda", run_config) + test_compact_matches_native_mtp(dtype, 4, run_config) else: test_replay_acceptance_flush_and_graph(3, (32, 64), 8, dtype, run_config) diff --git a/unit_tests/common/basemodel/triton_kernel/linear_att/test_replayssm_execution.py b/unit_tests/common/basemodel/triton_kernel/linear_att/test_replayssm_execution.py index dc2ddce177..99cf56f854 100644 --- a/unit_tests/common/basemodel/triton_kernel/linear_att/test_replayssm_execution.py +++ b/unit_tests/common/basemodel/triton_kernel/linear_att/test_replayssm_execution.py @@ -64,9 +64,9 @@ def metadata(cache): cu, reqs, _ = build_dynamic_mtp_linear_att_state_params(rows, mtp, manager.req_to_mtp_state_index, 4) return cu, reqs, cache.prepare_decode(reqs, cu) - def forward_accept(index, cu, reqs, pos): + def forward_accept(index, cu, reqs): cache = caches[index] - out = [cache.forward(layer, q, k, v, a, beta, alog, bias, reqs, pos, cu) for layer in range(3)] + out = [cache.forward(layer, q, k, v, a, beta, alog, bias, reqs, cu) for layer in range(3)] if index == 0: cache.accept_updates(reqs, accepted) else: @@ -77,7 +77,7 @@ def forward_accept(index, cu, reqs, pos): graphs, captured, outputs = [], [], [] for index, cache in enumerate(caches): cu, reqs, pos = metadata(cache) - forward_accept(index, cu, reqs, pos) + forward_accept(index, cu, reqs) cache.state.copy_(initial) cache.cursors.zero_() manager.req_to_mtp_state_index.zero_() @@ -85,7 +85,7 @@ def forward_accept(index, cu, reqs, pos): with torch.cuda.graph(g): if capture_metadata: cu, reqs, pos = metadata(cache) - out = forward_accept(index, cu, reqs, pos) + out = forward_accept(index, cu, reqs) graphs.append(g) captured.append((cu, reqs, pos)) outputs.append(out) @@ -147,12 +147,12 @@ def test_decode_accept_order_preserves_forward_snapshot(dtype, capture_metadata, alog = torch.zeros(4, device="cuda") bias = torch.zeros_like(alog) - def forward(cache, pos): - return [cache.forward(layer, q, k, v, a, beta, alog, bias, reqs, pos) for layer in range(3)] + def forward(cache): + return [cache.forward(layer, q, k, v, a, beta, alog, bias, reqs) for layer in range(3)] def warm(cache): - pos = cache.prepare_decode(reqs) - out = forward(cache, pos) + cache.prepare_decode(reqs) + out = forward(cache) cache.accept_updates(reqs) return out @@ -173,7 +173,7 @@ def warm(cache): pos = cache.prepare_decode(reqs) if early: cache.accept_updates(reqs) - out = forward(cache, pos) + out = forward(cache) if capture_metadata and not early: cache.accept_updates(reqs) graphs.append(graph) @@ -234,7 +234,7 @@ def test_cross_layer_fold_with_state_offset_above_int32(): def step(cache, accepted): pos = cache.prepare_decode(reqs, cu) - out = [cache.forward(layer, q, k, v, a, beta, alog, bias, reqs, pos, cu) for layer in range(48)] + out = [cache.forward(layer, q, k, v, a, beta, alog, bias, reqs, cu) for layer in range(48)] cache.accept_updates(reqs, accepted) return pos, out @@ -281,8 +281,8 @@ def test_small_batch_matches_larger_graph_padding(batch, dtype): for cache, ids, offsets in zip(caches, [reqs, padded], [cu, padded_cu]): def step(): - pos = cache.prepare_decode(ids, offsets) - out = [cache.forward(layer, q, k, v, a, beta, alog, bias, ids, pos, offsets) for layer in range(2)] + cache.prepare_decode(ids, offsets) + out = [cache.forward(layer, q, k, v, a, beta, alog, bias, ids, offsets) for layer in range(2)] cache.accept_updates(ids, accepted) return out diff --git a/unit_tests/common/basemodel/triton_kernel/linear_att/test_replayssm_grouped_keys.py b/unit_tests/common/basemodel/triton_kernel/linear_att/test_replayssm_grouped_keys.py index ff8e7de469..ea4a2795a7 100644 --- a/unit_tests/common/basemodel/triton_kernel/linear_att/test_replayssm_grouped_keys.py +++ b/unit_tests/common/basemodel/triton_kernel/linear_att/test_replayssm_grouped_keys.py @@ -9,7 +9,7 @@ pytestmark = pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") -@pytest.mark.parametrize("mode", ["replay", "compact", "kda"]) +@pytest.mark.parametrize("mode", ["replay", "compact"]) @pytest.mark.parametrize("group", [1, 2, 3]) @pytest.mark.parametrize("dtype", [torch.float32, torch.bfloat16]) def test_grouped_keys_graph_fold_partial_accept(mode, group, dtype): @@ -21,16 +21,14 @@ def test_grouped_keys_graph_fold_partial_accept(mode, group, dtype): def make_cache(key_heads): if mode == "replay": return ReplaySSMCache(state.clone(), 16, width, num_key_heads=key_heads) - return CompactSSMCache( - state.clone(), width, torch.bfloat16, kind="kda" if mode == "kda" else "gdn", num_key_heads=key_heads - ) + return CompactSSMCache(state.clone(), width, torch.bfloat16, num_key_heads=key_heads) caches = [make_cache(hv), make_cache(h)] assert caches[0].keys.numel() == caches[1].keys.numel() * group # Packed strides, out-of-order requests, an empty real request, nonempty HOLD. packed = torch.randn(1, 8, (2 * h + hv) * kd, device="cuda", dtype=torch.bfloat16) q, k, v = [x.view(1, 8, heads, kd) for x, heads in zip(packed.split([h * kd, h * kd, hv * kd], -1), [h, h, hv])] - a = torch.randn(8, hv * kd if mode == "kda" else hv, device="cuda", dtype=torch.bfloat16) + a = torch.randn(8, hv, device="cuda", dtype=torch.bfloat16) b = torch.randn(8, hv, device="cuda", dtype=torch.bfloat16) alog = torch.randn(hv, device="cuda") * 0.1 bias = torch.randn(a.shape[-1], device="cuda") * 0.1 @@ -39,8 +37,8 @@ def make_cache(key_heads): accepted = torch.tensor([0, -1, 0, -1, -1], device="cuda", dtype=torch.int32) def step(cache): - pos = cache.prepare_decode(reqs, cu) - out = [cache.forward(layer, q, k, v, a, b, alog, bias, reqs, pos, cu) for layer in range(layers)] + cache.prepare_decode(reqs, cu) + out = [cache.forward(layer, q, k, v, a, b, alog, bias, reqs, cu) for layer in range(layers)] cache.accept_updates(reqs, accepted) return out diff --git a/unit_tests/common/basemodel/triton_kernel/linear_att/test_ssm_autotune.py b/unit_tests/common/basemodel/triton_kernel/linear_att/test_ssm_autotune.py index 92174f8f25..199bdc4ae9 100644 --- a/unit_tests/common/basemodel/triton_kernel/linear_att/test_ssm_autotune.py +++ b/unit_tests/common/basemodel/triton_kernel/linear_att/test_ssm_autotune.py @@ -19,14 +19,14 @@ def make_inputs(mode, dtype=torch.bfloat16, device="cpu", width=4, projection_mo if mode == "replay": cache = ReplaySSMCache(state, 8, width, num_key_heads=h, projection_mode=projection_mode) else: - cache = CompactSSMCache(state, width, torch.bfloat16, mode, num_key_heads=h) + cache = CompactSSMCache(state, width, torch.bfloat16, num_key_heads=h) tokens = batch * width packed = torch.randn(1, tokens, 2 * h * kd + hv * vd, dtype=torch.bfloat16, device=device) q, k, v = [ x.view(1, tokens, heads, dim) for x, heads, dim in zip(packed.split([h * kd, h * kd, hv * vd], -1), [h, h, hv], [kd, kd, vd]) ] - gates = torch.randn(tokens, hv * (kd + 1 if mode == "kda" else 2), device=device, dtype=q.dtype) + gates = torch.randn(tokens, hv * 2, device=device, dtype=q.dtype) a, b = gates[:, :-hv], gates[:, -hv:] return dict( cache=cache, @@ -75,7 +75,7 @@ def isolate_tuner(monkeypatch, tmp_path): monkeypatch.setattr(tuner, "warmuped_configs_set", set()) -@pytest.mark.parametrize("mode", ["gdn", "kda", "replay"]) +@pytest.mark.parametrize("mode", ["gdn", "replay"]) @pytest.mark.parametrize("dtype", [torch.float32, torch.bfloat16]) def test_rebuild_has_real_requests_and_independent_packed_inputs(mode, dtype): inputs = make_inputs(mode, dtype) @@ -139,7 +139,6 @@ def test_layout_is_fixed_before_graph_capture_and_cached_configs_do_not_execute( "mode,width,projection_mode", [ ("gdn", 4, "inline"), - ("kda", 4, "inline"), ("replay", 4, "inline"), ("replay", 1, "inline"), ("replay", 4, "precompute"), @@ -175,10 +174,10 @@ def checked_bench(*args, **kwargs): monkeypatch.setattr(tuner, "_bench", checked_bench) args = [inputs[name] for name in ("q", "k", "v", "a", "b", "a_log", "bias")] reqs = torch.full((3,), cache.hold, device="cuda", dtype=torch.int32) - positions = cache.prepare_decode(reqs, inputs["cu_seqlens"]) + cache.prepare_decode(reqs, inputs["cu_seqlens"]) with Autotuner.autotune_warmup(AutotuneKernelType.DECODE_ATTENTION): # Exercise the real forward hook with the HOLD-only graph warmup input. - cache.forward(0, *args, reqs, positions, inputs["cu_seqlens"]) + cache.forward(0, *args, reqs, inputs["cu_seqlens"]) assert len(timings) == len(configs) assert cache.run_config in configs # An explicit identical layout supplies a reference for capture/replay and @@ -190,11 +189,8 @@ def checked_bench(*args, **kwargs): accepted = torch.zeros(cache.state.shape[1], device="cuda", dtype=torch.int32) def step(target): - positions = target.prepare_decode(reqs, inputs["cu_seqlens"]) - out = [ - target.forward(layer, *args, reqs, positions, inputs["cu_seqlens"]) - for layer in range(target.state.shape[0]) - ] + target.prepare_decode(reqs, inputs["cu_seqlens"]) + out = [target.forward(layer, *args, reqs, inputs["cu_seqlens"]) for layer in range(target.state.shape[0])] target.accept_updates(reqs, accepted) return out From 9f791ad56959dc7e7e94e3e44a757613292e0dde Mon Sep 17 00:00:00 2001 From: sufubao Date: Tue, 6 Oct 2026 13:37:28 +0800 Subject: [PATCH 05/40] feat: support ReplaySSM state modes for KDA --- docs/replayssm/README.md | 25 ++- .../common/basemodel/attention/linear/kda.py | 30 +++- .../triton_kernel/linear_att/replayssm.py | 159 +++++++++++++++--- .../linear_att/replayssm_compact.py | 37 +++- .../triton_kernel/linear_att/ssm_autotune.py | 15 +- lightllm/common/req_manager/glm5_next.py | 2 + lightllm/common/req_manager/linear_att.py | 4 + lightllm/server/api_cli.py | 2 +- lightllm/server/api_start.py | 4 +- test/benchmark/kernels/benchmark_kda_ssm.py | 148 ++++++++++++++++ .../linear_att/test_replayssm.py | 94 +++++++++-- .../linear_att/test_replayssm_execution.py | 48 ++++-- .../linear_att/test_ssm_autotune.py | 37 ++-- unit_tests/models/glm5_next/test_cache.py | 37 +++- unit_tests/models/glm5_next/test_mtp.py | 139 +++++++++++++-- unit_tests/models/glm5_next/test_pd_cache.py | 57 ++++++- unit_tests/server/test_ssm_state_mode.py | 3 + 17 files changed, 731 insertions(+), 110 deletions(-) create mode 100644 test/benchmark/kernels/benchmark_kda_ssm.py diff --git a/docs/replayssm/README.md b/docs/replayssm/README.md index f3a1a5a45c..1e9dd560a0 100644 --- a/docs/replayssm/README.md +++ b/docs/replayssm/README.md @@ -1,4 +1,4 @@ -# ReplaySSM for GDN +# ReplaySSM for GDN and KDA ReplaySSM reduces speculative SSM storage by keeping one checkpoint per request and replaying accepted updates. It is opt-in; `--ssm_state_mode native` remains @@ -25,10 +25,11 @@ capacities are 4, 8, 16, 32 and 64. | `compact` | One state plus this round's raw inputs; merge the accepted prefix every round. Requires MTP. | | `replay` | One state plus bounded accepted history; fold before the next verify window would overflow. Supports ordinary decode and MTP. | -Non-native modes currently support Qwen3-Next and Qwen3.5 GDN with FP32 or BF16 +Non-native modes support Qwen3-Next/Qwen3.5 GDN and GLM-5.3-FLASH KDA with FP32 or BF16 state. Other models retain their native path. `--replayssm_projection_mode` selects `inline` (default) or `precompute`; the latter is only valid with replay. Choose the projection mode and capacity using the target workload. +Without MTP, replay adds history storage to the single native checkpoint. ## State lifecycle and precision @@ -53,7 +54,15 @@ The SSM autotuner measures prepare, forward, acceptance and fold on disposable state. Its selected layout is fixed before CUDA Graph capture and reused for verification and checkpoint reconstruction; serving state is never a tuning input. Cached configurations distinguish precision, stride, heads, verify width, -history capacity and projection mode. +history capacity, projection mode and recurrent rule. + +KDA uses a per-key-dimension log decay, while GDN uses one scalar per value +head. KDA history stores the full gate vector; checkpoint projection applies +the cumulative vector decay to Q/K before projecting. The low-rank correction +uses forward decay between accepted tokens, without inverse exponential +scaling. The model's `gate_lower_bound` also separates autotune configurations. +Accepted-state exports retain GLM's indexer tail for PD; CPU prefix restore +clears the tail at aligned pool boundaries. ## Memory and performance @@ -63,6 +72,11 @@ request. Replay L8/precompute uses about 24.07 MiB including history and cursors Conv state, full-attention KV, request tables and sampling buffers are additional allocations. +For GLM-5.3-FLASH TP8 (34 KDA layers, 8 local heads), one FP32 checkpoint +is 17 MiB per GPU. Native MTP2 uses 51 MiB per request; replay L8/inline +uses about 25.52 MiB including history. With BF16 state the corresponding +figures are 25.5 and 17.02 MiB. + More KV space improves throughput only when the workload and request capacity can use it. Increasing request capacity can also increase per-token latency. Keep algorithm comparisons at equal capacity separate from capacity comparisons. @@ -70,14 +84,15 @@ Keep algorithm comparisons at equal capacity separate from capacity comparisons. ## Tests ```bash -PYTHONPATH=. python -m pytest -q \ +exp -m "ReplaySSM GDN/KDA lifecycle regression" /usr/bin/env PYTHONPATH=. python -m pytest -q \ unit_tests/server/test_ssm_state_mode.py \ unit_tests/common/basemodel/attention/linear/test_gdn.py \ unit_tests/common/basemodel/triton_kernel/linear_att/test_replayssm*.py \ unit_tests/common/basemodel/triton_kernel/linear_att/test_ssm_autotune.py \ unit_tests/common/basemodel/triton_kernel/linear_att/test_mtp_state_params.py \ unit_tests/common/basemodel/triton_kernel/linear_att/test_acceptance_fusion.py \ - unit_tests/common/test_paged_kv_transfer.py + unit_tests/common/test_paged_kv_transfer.py \ + unit_tests/models/glm5_next/test_{mtp,cache,pd_cache}.py ``` CUDA tests cover partial acceptance, history folds, BF16 rounding, grouped keys, diff --git a/lightllm/common/basemodel/attention/linear/kda.py b/lightllm/common/basemodel/attention/linear/kda.py index 367d9b6c9e..a90c0e42ef 100644 --- a/lightllm/common/basemodel/attention/linear/kda.py +++ b/lightllm/common/basemodel/attention/linear/kda.py @@ -45,6 +45,8 @@ def __init__(self, model: "TpPartBaseModel"): self.conv_kernel_size = config["short_conv_kernel_size"] self.lower_bound = config.get("gate_lower_bound", -5.0) self.mtp_step = model.args.mtp_step + if model.req_manager.ssm_update_cache is not None: + model.req_manager.ssm_update_cache.lower_bound = self.lower_bound def create_att_prefill_state(self, infer_state: "InferStateInfo"): return KDAPrefillAttState(backend=self, infer_state=infer_state) @@ -63,7 +65,10 @@ class KDAPrefillAttState(BasePrefillAttState): def init_state(self): self.b_conv_buffer_idx = self.infer_state.b_req_idx - self.b_ssm_buffer_idx = self.infer_state.b_req_idx * (self.backend.mtp_step + 1) + manager = self.infer_state.req_manager + self.b_ssm_buffer_idx = self.infer_state.b_req_idx * manager.ssm_slots_per_req + if manager.ssm_update_cache is not None: + manager.ssm_update_cache.merge_accepted_updates(self.infer_state.b_req_idx) def prefill_att( self, @@ -137,6 +142,11 @@ def init_state(self): self._init_dynamic_mtp_decode_state(mtp_step + 1) else: self._init_fixed_mtp_decode_state(mtp_step) + updates = self.infer_state.req_manager.ssm_update_cache + if updates is not None: + updates.prepare_decode(self.b_conv_buffer_idx, self.b1_mtp_cu_q_seq_len) + if mtp_step == 0: + updates.accept_updates(self.b_conv_buffer_idx) def _init_normal_decode_state(self): self.b_conv_buffer_idx = self.infer_state.b_req_idx @@ -176,6 +186,9 @@ def _init_fixed_mtp_decode_state(self, mtp_step: int): self._init_mtp_ssm_buffer_idx(mtp_size) def _init_mtp_ssm_buffer_idx(self, mtp_size: int): + if self.infer_state.req_manager.ssm_update_cache 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) @@ -221,6 +234,21 @@ def decode_att( v = v.view(*shape, backend.tp_num_heads, backend.head_dim) raw_gate = raw_gate.view(*shape, backend.tp_hidden_size) raw_beta = raw_beta.view(*shape, backend.tp_num_heads) + updates = self.infer_state.req_manager.ssm_update_cache + if updates is not None: + layer = layer_num - layer_num // self.infer_state.req_manager.linear_config.full_attention_interval + return updates.forward( + layer, + q, + k, + v, + raw_gate.view(-1, backend.tp_hidden_size), + raw_beta.view(-1, backend.tp_num_heads), + layer_weight.linear_A_log.weight, + layer_weight.linear_dt_bias.weight, + self.b_conv_buffer_idx, + self.b1_mtp_cu_q_seq_len, + ) output, _ = fused_recurrent_kda( q=q, k=k, diff --git a/lightllm/common/basemodel/triton_kernel/linear_att/replayssm.py b/lightllm/common/basemodel/triton_kernel/linear_att/replayssm.py index 74ce36492e..d68eedb22d 100644 --- a/lightllm/common/basemodel/triton_kernel/linear_att/replayssm.py +++ b/lightllm/common/basemodel/triton_kernel/linear_att/replayssm.py @@ -1,5 +1,5 @@ # SPDX-License-Identifier: Apache-2.0 -"""GDN ReplaySSM in LightLLM's [K, V] state layout. +"""GDN/KDA 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 @@ -101,6 +101,8 @@ def _fold_history( RS: tl.constexpr, VS: tl.constexpr, GS: tl.constexpr, + BS: tl.constexpr, + KDA: tl.constexpr, ): count = tl.load(Active + N) work = tl.program_id(0) @@ -125,13 +127,21 @@ def _fold_history( for j in range(n): raw_k = tl.load(RawKeys + layer * RS + (key_slot * (2 * L) + base + j) * K + kk, kk < K, 0).to(tl.float32) raw_v = tl.load(RawValues + layer * VS + (slot * (2 * L) + base + j) * V + vv, vv < V, 0).to(tl.float32) - g = tl.load(Gates + layer * GS + slot * (2 * L) + base + j).to(tl.float32) - beta = tl.load(Betas + layer * GS + slot * (2 * L) + base + j).to(tl.float32) + if KDA: + g = tl.load(Gates + layer * GS + (slot * (2 * L) + base + j) * K + kk, kk < K, 0).to(tl.float32) + else: + g = tl.load(Gates + layer * GS + slot * (2 * L) + base + j).to(tl.float32) + beta = tl.load(Betas + layer * BS + slot * (2 * L) + base + j).to(tl.float32) raw_k /= tl.sqrt(tl.sum(raw_k * raw_k) + 1.0e-6) decay = tl.exp(g) - sk = tl.sum(state * raw_k[:, None], 0) - d = beta * (raw_v - decay * sk) - state = state * decay + raw_k[:, None] * d[None, :] + if KDA: + state *= decay[:, None] + d = beta * (raw_v - tl.sum(state * raw_k[:, None], 0)) + state += raw_k[:, None] * d[None, :] + else: + sk = tl.sum(state * raw_k[:, None], 0) + d = beta * (raw_v - decay * sk) + state = state * decay + raw_k[:, None] * d[None, :] # Start the new history from the same rounded checkpoint future calls load. state = state.to(State.dtype.element_ty).to(tl.float32) tl.store(sp, state, (kk[:, None] < K) & (vv[None, :] < V)) @@ -177,6 +187,8 @@ def _replay( PRECOMPUTE_STATE: tl.constexpr, BW: tl.constexpr, FUSED_FOLD: tl.constexpr, + KDA: tl.constexpr, + LOWER: tl.constexpr, ): """Output reconstruction with an exact raw-input recurrence at checkpoint folds.""" iv, seq, hv = tl.program_id(0), tl.program_id(1), tl.program_id(2) @@ -207,13 +219,21 @@ def _replay( for j in range(n): raw_k = tl.load(RawKeys + (key_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) + if KDA: + g = tl.load(Gates + (slot * (2 * L) + base + j) * K + kk, kk < K, 0).to(tl.float32) + else: + 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) decay = tl.exp(g) - sk = tl.sum(state * raw_k[:, None], 0) - d = beta * (raw_v - decay * sk) - state = state * decay + raw_k[:, None] * d[None, :] + if KDA: + state *= decay[:, None] + d = beta * (raw_v - tl.sum(state * raw_k[:, None], 0)) + state += raw_k[:, None] * d[None, :] + else: + sk = tl.sum(state * raw_k[:, None], 0) + d = beta * (raw_v - decay * sk) + state = state * decay + raw_k[:, None] * d[None, :] # Start the new history from the same rounded checkpoint future calls load. state = state.to(State.dtype.element_ty).to(tl.float32) tl.store(sp, state, (kk[:, None] < K) & (vv[None, :] < V)) @@ -241,8 +261,41 @@ def _replay( qk = q_rows + k_rows qk *= tl.rsqrt(tl.sum(qk * qk, axis=1) + 1.0e-6)[:, None] qk *= tl.where(tt < BW, K ** -0.5, 1.0)[:, None] - projections = tl.dot(qk.to(projection_dtype), state.to(projection_dtype), input_precision=projection_precision) + if KDA: + history = ll < n + gates = tl.load( + Gates + (slot * (2 * L) + base + ll[:, None]) * K + kk[None, :], + history[:, None] & (kk[None, :] < K), + 0, + ).to(tl.float32) + total_g = tl.sum(gates, axis=0) + wi = tl.arange(0, BW) + gate_inputs = tl.load( + A + (start + wi[:, None]) * SA + hv * K + kk[None, :], + (start + wi[:, None] < end) & (kk[None, :] < K), + 0, + ).to(tl.float32) + gate_inputs += tl.load(Bias + hv * K + kk, kk < K, 0).to(tl.float32)[None, :] + fresh_gates = LOWER * tl.sigmoid(tl.exp(tl.load(Alog + hv)) * gate_inputs) + fresh_gates = tl.where((start + wi < end)[:, None], fresh_gates, 0) + query_gates = total_g[None, :] + tl.gather( + tl.cumsum(fresh_gates, axis=0), + tl.broadcast_to((tt % BW)[:, None], qk.shape), + 0, + ) + # All exponents are forward decay. No inverse prefix scaling can + # overflow when a long history contains strongly negative gates. + projections = tl.dot( + (qk * tl.exp(query_gates)).to(projection_dtype), + state.to(projection_dtype), + input_precision=projection_precision, + ) + else: + projections = tl.dot( + qk.to(projection_dtype), state.to(projection_dtype), input_precision=projection_precision + ) + if PRECOMPUTE_STATE and not KDA: # Project the whole verify window against both accepted and new keys. # Only the small corrected-value solve below depends on token order. jj = tl.arange(0, triton.next_power_of_2(max(16, L + BW))) @@ -328,7 +381,7 @@ def _replay( tl.store(Betas + slot * (2 * L) + records, betas, fresh) return - else: + if not PRECOMPUTE_STATE or KDA: history = ll < n keys = tl.load( Keys + (key_slot * (2 * L) + base) * K + ll[:, None] * K + kk[None, :], @@ -340,11 +393,21 @@ def _replay( history[:, None] & (vv[None, :] < V), 0, ).to(tl.float32) - gates = tl.load(Gates + slot * (2 * L) + base + ll, history, 0).to(tl.float32) + if KDA: + gates = tl.load( + Gates + (slot * (2 * L) + base + ll[:, None]) * K + kk[None, :], + history[:, None] & (kk[None, :] < K), + 0, + ).to(tl.float32) + else: + 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) + if KDA: + keys *= tl.where(history[:, None], tl.exp(total_g[None, :] - gate_prefix), 0.0) h = hv // (HV // H) - log_a, bias = tl.load(Alog + hv).to(tl.float32), tl.load(Bias + hv).to(tl.float32) + log_a = tl.load(Alog + hv).to(tl.float32) + bias = tl.load(Bias + hv * K + kk, kk < K, 0).to(tl.float32) if KDA else 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) @@ -352,17 +415,36 @@ def _replay( 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) + if KDA: + x = tl.load(A + t * SA + hv * K + kk, kk < K, 0).to(tl.float32) + bias + g = LOWER * tl.sigmoid(tl.exp(log_a) * x) + else: + 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 - checkpoint_sk = tl.sum(state * k[:, None], axis=0) - checkpoint_sq = tl.sum(state * q[:, None], axis=0) - sk = checkpoint_sk * tl.exp(total_g) + tl.sum(deltas * hk[:, None], axis=0) - sq = checkpoint_sq * tl.exp(total_g) + tl.sum(deltas * hq[:, None], axis=0) + if KDA: + # Advance the decayed keys once per token, rather than + # evaluating L independent vector exponentials each time. + keys *= tl.exp(g)[None, :] + hk = tl.sum(keys * k[None, :], axis=1) + hq = tl.sum(keys * q[None, :], axis=1) + if PRECOMPUTE_STATE: + checkpoint_sk = tl.sum(tl.where((tt == BW + t - start)[:, None], projections, 0), 0) + checkpoint_sq = tl.sum(tl.where((tt == t - start)[:, None], projections, 0), 0) + else: + checkpoint_sk = tl.sum(state * (k * tl.exp(total_g))[:, None], axis=0) + checkpoint_sq = tl.sum(state * (q * tl.exp(total_g))[:, None], axis=0) + sk = checkpoint_sk + tl.sum(deltas * hk[:, None], axis=0) + sq = checkpoint_sq + tl.sum(deltas * hq[:, None], axis=0) + else: + 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 + checkpoint_sk = tl.sum(state * k[:, None], axis=0) + checkpoint_sq = tl.sum(state * q[:, None], axis=0) + sk = checkpoint_sk * tl.exp(total_g) + tl.sum(deltas * hk[:, None], axis=0) + sq = checkpoint_sq * 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 @@ -374,11 +456,15 @@ def _replay( if hv % (HV // KH) == 0: tl.store(Keys + (key_slot * (2 * L) + record) * K + kk, k, kk < K) tl.store(RawKeys + (key_slot * (2 * L) + record) * K + kk, raw_k, kk < K) - tl.store(Gates + slot * (2 * L) + record, g) + if KDA: + tl.store(Gates + (slot * (2 * L) + record) * K + kk, g, kk < K) + else: + 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) + if not KDA: + gate_prefix = tl.where(ll >= n, total_g, gate_prefix) history |= ll == n n += 1 @@ -403,6 +489,7 @@ def _materialize( BK: tl.constexpr, BV: tl.constexpr, SNAPSHOT: tl.constexpr, + KDA: 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) @@ -423,10 +510,13 @@ def _materialize( for j in range(n): raw_k = tl.load(RawKeys + (key_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) + if KDA: + g = tl.load(Gates + (slot * (2 * L) + base + j) * K + kk, kk < K, 0).to(tl.float32) + else: + 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) + state *= tl.exp(g)[:, None] if KDA else tl.exp(g) d = beta * (raw_v - tl.sum(state * raw_k[:, None], 0)) state += raw_k[:, None] * d[None, :] out_slot = layer * HV + head if SNAPSHOT else slot @@ -445,6 +535,8 @@ def __init__( activation_dtype=torch.bfloat16, *, num_key_heads=None, + kda=False, + lower_bound=-5.0, projection_mode="inline", run_config=None, ): @@ -452,6 +544,8 @@ def __init__( assert verify_width > 0 and capacity >= max(4, verify_width) and capacity & (capacity - 1) == 0 assert projection_mode in ("inline", "precompute") self.state = state + self.kda = kda + self.lower_bound = lower_bound self.fold_programs = ( 8 * torch.cuda.get_device_properties(state.device).multi_processor_count if state.is_cuda else 0 ) @@ -477,7 +571,11 @@ def __init__( key_shape = (layers, slots, self.num_key_heads, 2 * capacity, k) self.keys = torch.empty(key_shape, dtype=history_dtype, device=state.device) self.deltas = torch.empty((layers, slots, hv, 2 * capacity, v), dtype=history_dtype, device=state.device) - self.gates = torch.empty((layers, slots, hv, 2 * capacity), dtype=torch.float32, device=state.device) + self.gates = torch.empty( + (layers, slots, hv, 2 * capacity, k) if kda else (layers, slots, hv, 2 * capacity), + dtype=torch.float32, + device=state.device, + ) self.raw_keys = torch.empty(key_shape, 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) @@ -536,6 +634,8 @@ def prepare_decode(self, reqs, cu_seqlens=None): self.raw_keys.stride(0), self.raw_values.stride(0), self.gates.stride(0), + self.betas.stride(0), + self.kda, num_warps=2, num_stages=1, ) @@ -583,6 +683,7 @@ def _materialize_accepted_state(self, reqs, output, snapshot): triton.next_power_of_2(k), bv, snapshot, + self.kda, num_warps=config["num_warps"], num_stages=config.get("num_stages", 3), ) @@ -645,6 +746,8 @@ def forward(self, layer, q, k, v, a, b, a_log, bias, reqs, cu_seqlens=None, run_ precompute_state, triton.next_power_of_2(self.verify_width), reqs.numel() > 12, + self.kda, + self.lower_bound, num_warps=config["num_warps"], num_stages=config.get("num_stages", 3), ) diff --git a/lightllm/common/basemodel/triton_kernel/linear_att/replayssm_compact.py b/lightllm/common/basemodel/triton_kernel/linear_att/replayssm_compact.py index 81e19c6d91..6799be5f34 100644 --- a/lightllm/common/basemodel/triton_kernel/linear_att/replayssm_compact.py +++ b/lightllm/common/basemodel/triton_kernel/linear_att/replayssm_compact.py @@ -1,5 +1,5 @@ # SPDX-License-Identifier: Apache-2.0 -"""Compact speculative state for GDN. +"""Compact speculative state for GDN 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. @@ -46,6 +46,8 @@ def _compact( 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) @@ -72,23 +74,34 @@ def _compact( if COMMIT: k = tl.load(Keys + key_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) + if KDA: + decay = tl.load(Decays + record * K + kk, kk < K, 1)[:, None] + else: + decay = 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) - 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) + if KDA: + x = tl.load(A + t * SA + hv * K + kk, kk < K, 0).to(tl.float32) + x += tl.load(Bias + hv * K + kk, kk < K, 0).to(tl.float32) + decay = tl.exp(LOWER * tl.sigmoid(tl.exp(log_a) * x))[:, None] + 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: if hv % (HV // KH) == 0: tl.store(Keys + key_record * K + kk, k, kk < K) tl.store(Betas + record, beta) - tl.store(Decays + record, decay) + if KDA: + tl.store(Decays + record * K + kk, tl.sum(decay, 1), kk < K) + else: + tl.store(Decays + record, decay) k = k / tl.sqrt(tl.sum(k * k) + 1.0e-6) state *= decay d = (v - tl.sum(state * k[:, None], 0)) * beta @@ -103,8 +116,12 @@ def _compact( class CompactSSMCache: - def __init__(self, state, verify_width, activation_dtype, *, num_key_heads=None, run_config=None): + def __init__( + self, state, verify_width, activation_dtype, *, num_key_heads=None, kda=False, lower_bound=-5.0, run_config=None + ): self.state = state + self.kda = kda + self.lower_bound = lower_bound self.verify_width = verify_width # Verify and accepted-state reconstruction must share the same layout. self._config_is_fixed = run_config is not None @@ -120,7 +137,7 @@ def __init__(self, state, verify_width, activation_dtype, *, num_key_heads=None, (layers, slots, self.num_key_heads, verify_width, k), device=state.device, dtype=activation_dtype ) self.values = torch.empty((*shape, v), device=state.device, dtype=activation_dtype) - self.decays = torch.empty(shape, device=state.device, dtype=torch.float32) + self.decays = torch.empty((*shape, k) if kda else shape, device=state.device, dtype=torch.float32) self.betas = torch.empty(shape, device=state.device, dtype=torch.float32) def clear_history(self, req): @@ -180,6 +197,8 @@ def forward(self, layer, q, k, v, a, b, a_log, bias, reqs, cu_seqlens=None): triton.next_power_of_2(kd), bv, False, + self.kda, + self.lower_bound, num_warps=self.run_config["num_warps"], num_stages=self.run_config.get("num_stages", 3), ) @@ -222,6 +241,8 @@ def accept_updates(self, reqs, accepted): triton.next_power_of_2(kd), bv, True, + self.kda, + self.lower_bound, num_warps=self.run_config["num_warps"], num_stages=self.run_config.get("num_stages", 3), ) diff --git a/lightllm/common/basemodel/triton_kernel/linear_att/ssm_autotune.py b/lightllm/common/basemodel/triton_kernel/linear_att/ssm_autotune.py index 6bb435cb77..d1b3fac568 100644 --- a/lightllm/common/basemodel/triton_kernel/linear_att/ssm_autotune.py +++ b/lightllm/common/basemodel/triton_kernel/linear_att/ssm_autotune.py @@ -28,9 +28,11 @@ def static_key(cache, mode, q, k, v, a, b, cu_seqlens): "dtype": str((q.dtype, cache.state.dtype)), "strides": str((q.stride(axis), k.stride(axis), v.stride(axis), a.stride(0), b.stride(0))), "history": cache.capacity if mode == "replay" else 0, - "gate": 0, + "gate": int(cache.kda), "varlen": cu_seqlens is not None, } + if cache.kda: + key["lower"] = cache.lower_bound if mode == "replay": key["history_dtype"] = str(cache.keys.dtype) key["v"] = 4 # Small batches omit the forward fold already handled by prepare. @@ -58,10 +60,19 @@ def rebuild_inputs(cache, mode, q, k, v, a, b, a_log, bias, cu_seqlens=None, wor cache.verify_width, q.dtype, num_key_heads=cache.num_key_heads, + kda=cache.kda, + lower_bound=cache.lower_bound, projection_mode=cache.projection_mode, ) else: - scratch = type(cache)(state, cache.verify_width, q.dtype, num_key_heads=cache.num_key_heads) + scratch = type(cache)( + state, + cache.verify_width, + q.dtype, + num_key_heads=cache.num_key_heads, + kda=cache.kda, + lower_bound=cache.lower_bound, + ) # Benchmark callbacks must not recursively select a configuration. scratch._config_is_fixed = True inputs = [ diff --git a/lightllm/common/req_manager/glm5_next.py b/lightllm/common/req_manager/glm5_next.py index b4814afbf7..5814404b65 100644 --- a/lightllm/common/req_manager/glm5_next.py +++ b/lightllm/common/req_manager/glm5_next.py @@ -7,6 +7,8 @@ class Glm5NextReqManager(ReqManagerForMamba): """KDA runtime state and the NSA indexer's incomplete four-token pool.""" + kda = True + def __init__(self, max_request_num, max_sequence_length, mem_manager, linear_config): # Both checkpoint sizes are multiples of this hash page. A restored # prefix therefore has no incomplete K-pool to serialize or replay. diff --git a/lightllm/common/req_manager/linear_att.py b/lightllm/common/req_manager/linear_att.py index e8e1ab5170..8c796daa84 100644 --- a/lightllm/common/req_manager/linear_att.py +++ b/lightllm/common/req_manager/linear_att.py @@ -13,6 +13,8 @@ class ReqManagerForMamba(HybridAttentionReqManager): + kda = False + def __init__( self, max_request_num, @@ -64,6 +66,7 @@ def __init__( self.mtp_step + 1, linear_config.conv_state_dtype, num_key_heads=linear_config.num_linear_k_heads, + kda=self.kda, projection_mode=args.replayssm_projection_mode, ) elif args.ssm_state_mode == "compact": @@ -74,6 +77,7 @@ def __init__( self.mtp_step + 1, linear_config.conv_state_dtype, num_key_heads=linear_config.num_linear_k_heads, + kda=self.kda, ) return diff --git a/lightllm/server/api_cli.py b/lightllm/server/api_cli.py index dc1bd57155..b0d970e5c7 100644 --- a/lightllm/server/api_cli.py +++ b/lightllm/server/api_cli.py @@ -944,7 +944,7 @@ def add_cli_args(parser: argparse.ArgumentParser) -> argparse.ArgumentParser: "--ssm_state_mode", choices=["native", "compact", "replay"], default="native", - help="SSM state storage: native snapshots, compact MTP records, or deferred GDN replay (FP32/BF16).", + help="SSM state storage: native snapshots, compact MTP records, or deferred GDN/KDA replay (FP32/BF16).", ) parser.add_argument( "--replayssm_cache_len", diff --git a/lightllm/server/api_start.py b/lightllm/server/api_start.py index 3f05f1762b..8512baa3fc 100644 --- a/lightllm/server/api_start.py +++ b/lightllm/server/api_start.py @@ -49,12 +49,14 @@ def _validate_ssm_state_mode(args: StartArgs): from lightllm.utils.config_utils import get_model_type assert get_model_type(args.model_dir) in ( + "glm5_next", + "glm5_next_text", "qwen3_next", "qwen3_5", "qwen3_5_moe", "qwen3_5_text", "qwen3_5_moe_text", - ), "Non-native SSM state modes currently require a GDN model" + ), "Non-native SSM state modes currently require a GDN or KDA model" if args.ssm_state_mode == "compact": assert args.mtp_step > 0, "ssm_state_mode=compact requires MTP (mtp_step > 0)" else: diff --git a/test/benchmark/kernels/benchmark_kda_ssm.py b/test/benchmark/kernels/benchmark_kda_ssm.py new file mode 100644 index 0000000000..4a166228e2 --- /dev/null +++ b/test/benchmark/kernels/benchmark_kda_ssm.py @@ -0,0 +1,148 @@ +"""Measure a full KDA state cycle: all layers, acceptance, folds and export. + +This synthetic kernel benchmark excludes convolutions, model projections and +HTTP scheduling. Use a matched service workload for end-to-end conclusions. +""" + +import argparse +import gc +import json +import statistics +from pathlib import Path + +import torch +import triton + +from lightllm.common.basemodel.triton_kernel.linear_att.fla.ops.kda_decode import fused_recurrent_kda +from lightllm.common.basemodel.triton_kernel.linear_att.replayssm import ReplaySSMCache +from lightllm.common.basemodel.triton_kernel.linear_att.replayssm_compact import CompactSSMCache +from lightllm.common.basemodel.triton_kernel.linear_att.ssm_autotune import get_configs + + +def workload(args, batch, width, dtype, mode, acceptance): + layers, heads, dim = args.layers, args.heads, 128 + shape = (1, batch * width, heads, dim) if width > 1 else (batch, 1, heads, dim) + packed = torch.randn(batch * width, heads * dim * 3, dtype=torch.bfloat16, device="cuda") + q, k, v = [x.view(shape) for x in packed.split(heads * dim, -1)] + gate = torch.randn(batch * width, heads * dim, dtype=q.dtype, device="cuda") - 3 + beta = torch.randn(batch * width, heads, dtype=q.dtype, device="cuda") + alog = torch.zeros(heads, device="cuda") + bias = torch.zeros(heads * dim, device="cuda") + reqs = torch.arange(batch, dtype=torch.int32, device="cuda") + cu = torch.arange(batch + 1, dtype=torch.int32, device="cuda") * width if width > 1 else None + counts = reqs % width + 1 if acceptance == "partial" else torch.full_like(reqs, width) + accepted = torch.cat((counts - 1, counts.new_zeros(1))) + snapshot = torch.empty(layers, heads, dim, dim, dtype=dtype, device="cuda") + slots = (batch + 1) * (width if mode == "native" else 1) + state = torch.randn(layers, slots, heads, dim, dim, dtype=dtype, device="cuda") * 0.01 + cache = None + if mode == "compact": + cache = CompactSSMCache(state, width, q.dtype, kda=True) + elif mode != "native": + cache = ReplaySSMCache(state, 8, width, q.dtype, kda=True, projection_mode=mode) + if cache is not None: + cache._config_is_fixed = True + indexes = reqs[:, None] * width + torch.arange(width, device="cuda", dtype=torch.int32) + if width == 1: + indexes = reqs + bytes_used = ( + sum(x.numel() * x.element_size() for x in vars(cache).values() if isinstance(x, torch.Tensor)) + if cache + else state.numel() * state.element_size() + ) + + native_index = 0 if acceptance == "partial" else width - 1 + + def cycle(): + for _ in range(args.rounds): + if cache is not None: + cache.prepare_decode(reqs, cu) + for layer in range(layers): + if cache is None: + fused_recurrent_kda( + q, + k, + v, + gate.view(*shape[:2], -1), + beta.view(*shape[:2], heads), + alog, + bias, + state[layer], + indexes, + cu_seqlens=cu, + num_accepted_tokens=counts if cu is not None else None, + ) + else: + cache.forward(layer, q, k, v, gate, beta, alog, bias, reqs, cu) + if cache is not None: + cache.accept_updates(reqs, accepted) + if isinstance(cache, ReplaySSMCache): + cache._materialize_accepted_state(reqs[:1], snapshot, snapshot=True) + cache.merge_accepted_updates(reqs) + else: + snapshot.copy_(state[:, native_index] if cache is None else state[:, 0]) + + return cache, cycle, bytes_used + + +def main(): + p = argparse.ArgumentParser(description=__doc__) + p.add_argument("--batch", type=int, nargs="+", default=[16, 64, 192]) + p.add_argument("--width", type=int, nargs="+", default=[1, 3]) + p.add_argument("--dtype", nargs="+", choices=["float32", "bfloat16"], default=["float32", "bfloat16"]) + p.add_argument("--layers", type=int, default=34) + p.add_argument("--heads", type=int, default=8) + p.add_argument("--rounds", type=int, default=8) + p.add_argument("--repeats", type=int, default=5) + p.add_argument("--tune", action="store_true") + p.add_argument("--output", type=Path, required=True) + args = p.parse_args() + torch.manual_seed(53) + rows = [] + report = dict( + settings={**vars(args), "output": str(args.output)}, + device=torch.cuda.get_device_name(), + torch=torch.__version__, + triton=triton.__version__, + results=rows, + ) + for dtype_name in args.dtype: + for batch in args.batch: + for width in args.width: + for acceptance in ["full", "partial"] if width > 1 else ["full"]: + for mode in ["native", "compact", "inline", "precompute"] if width > 1 else ["native", "inline"]: + cache, cycle, size = workload(args, batch, width, getattr(torch, dtype_name), mode, acceptance) + candidates = [] + if cache is not None: + for config in get_configs() if args.tune else [{"BV": 32, "num_warps": 2, "num_stages": 1}]: + cache.run_config = config + cycle() + ms = triton.testing.do_bench_cudagraph(cycle, rep=50) + candidates.append(dict(config=config, ms=ms)) + cache.run_config = min(candidates, key=lambda r: r["ms"])["config"] + cycle() + samples = [ + triton.testing.do_bench_cudagraph(cycle, rep=100) / args.rounds for _ in range(args.repeats) + ] + row = dict( + batch=batch, + width=width, + dtype=dtype_name, + acceptance=acceptance, + mode=mode, + state_history_bytes=size, + ms_per_round=statistics.median(samples), + samples_ms=samples, + config=cache.run_config if cache else None, + candidates=candidates, + ) + rows.append(row) + args.output.write_text(json.dumps(report, indent=2) + "\n") + print(json.dumps(row), flush=True) + del cache, cycle + gc.collect() + torch.cuda.empty_cache() + + +if __name__ == "__main__": + main() 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 index a224ed0400..61dfbf791f 100644 --- a/unit_tests/common/basemodel/triton_kernel/linear_att/test_replayssm.py +++ b/unit_tests/common/basemodel/triton_kernel/linear_att/test_replayssm.py @@ -7,14 +7,18 @@ pytestmark = pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") -def reference(q, k, v, a, b, a_log, bias, state): +def reference(q, k, v, a, b, a_log, bias, state, kda=False, lower_bound=-5.0): 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] + if kda: + g = lower_bound * torch.sigmoid(a_log.exp()[:, None] * (a.float().view_as(k) + bias.view_as(k))) + state = state * g.exp()[..., None] + else: + 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 @@ -34,6 +38,8 @@ def test_replay_acceptance_flush_and_graph( activation_dtype=torch.bfloat16, gate_shift=-3, head_geometry=(2, 4), + kda=False, + lower_bound=-5.0, ): torch.manual_seed(21) kdim, vdim = dims @@ -45,6 +51,8 @@ def test_replay_acceptance_flush_and_graph( capacity, width, num_key_heads=h, + kda=kda, + lower_bound=lower_bound, projection_mode=projection_mode, run_config=run_config, activation_dtype=activation_dtype, @@ -63,10 +71,10 @@ def test_replay_acceptance_flush_and_graph( q = torch.randn(1, tokens, h, kdim, device="cuda", dtype=activation_dtype) k = torch.randn_like(q) v = torch.randn(1, tokens, hv, vdim, device="cuda", dtype=activation_dtype) - a = torch.randn(tokens, hv, device="cuda", dtype=activation_dtype) + gate_shift - b = torch.randn_like(a) + a = torch.randn(tokens, hv * kdim if kda else hv, device="cuda", dtype=activation_dtype) + gate_shift + b = torch.randn(tokens, hv, device="cuda", dtype=activation_dtype) alog = torch.randn(hv, device="cuda") * 0.1 - bias = torch.randn(hv, device="cuda") * 0.1 + bias = torch.randn(hv * kdim if kda else hv, device="cuda") * 0.1 accepted = torch.zeros(slots, dtype=torch.int32, device="cuda") def step(): @@ -103,7 +111,9 @@ def step(): 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) + out, cur = reference( + q[0, t], k[0, t], v[0, t], a[t], b[t], alog, bias, cur, kda=kda, lower_bound=lower_bound + ) 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) @@ -362,23 +372,23 @@ def test_snapshot_without_history_copies_active_state(): @pytest.mark.parametrize("dtype", [torch.float32, torch.bfloat16]) -def test_compact_preserves_accepted_prefix(dtype, run_config=None): +def test_compact_preserves_accepted_prefix(dtype, run_config=None, kda=False): from lightllm.common.basemodel.triton_kernel.linear_att.replayssm_compact import CompactSSMCache torch.manual_seed(17) - hv, kd, vd, width = 4, 64, 64, 3 + hv, kd, vd, width = 4, 128 if kda else 64, 128 if kda else 64, 3 state = torch.randn(2, 4, hv, kd, vd, device="cuda", dtype=dtype) * 0.01 - cache = CompactSSMCache(state, width, torch.bfloat16, num_key_heads=hv, run_config=run_config) + cache = CompactSSMCache(state, width, torch.bfloat16, num_key_heads=hv, kda=kda, run_config=run_config) 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), device="cuda", dtype=torch.bfloat16) + a = torch.randn((5, hv * kd if 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, device="cuda") * 0.1 + bias = torch.randn(hv * kd if 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, cu) for layer in range(2)] @@ -389,7 +399,7 @@ def test_compact_preserves_accepted_prefix(dtype, run_config=None): cur = before[layer, req].float() 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) + out, cur = reference(q[0, t], k[0, t], v[0, t], a[t], b[t], alog, bias, cur, kda=kda) 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]): @@ -448,8 +458,15 @@ def test_compact_matches_native_mtp(dtype, width, run_config=None): run_config={"num_stages": 1, **(run_config or {"BV": 8, "num_warps": 1})}, ) cache.accept_updates(reqs, accepted) - assert torch.equal(out, ref) - torch.testing.assert_close(state[0, :batch], reference_state[reqs * width + accepted[:batch]], rtol=0, atol=0) + # BF16 casts can straddle a rounding midpoint across compiler/layouts. + rtol = 2 * torch.finfo(dtype).eps if dtype == torch.bfloat16 else 0 + torch.testing.assert_close(out, ref, rtol=rtol, atol=2e-5 if dtype == torch.bfloat16 else 0) + torch.testing.assert_close( + state[0, :batch], + reference_state[reqs * width + accepted[:batch]], + rtol=rtol, + atol=5e-4 if dtype == torch.bfloat16 else 0, + ) counts.copy_(accepted[:batch] + 1) accepted[:batch].add_(1).remainder_(width) @@ -585,3 +602,50 @@ def test_checkpoint_restores_pending_history_and_accepted_conv(save_big, monkeyp 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 + + +@pytest.mark.parametrize("dtype", [torch.float32, torch.bfloat16]) +@pytest.mark.parametrize("capacity,width", [(4, 1), (4, 3), (8, 4), (16, 5), (32, 4), (64, 4)]) +@pytest.mark.parametrize("projection_mode", ["inline", "precompute"]) +@pytest.mark.parametrize("lower_bound,gate_shift", [(-5.0, 0), (-2.5, -6)]) +@pytest.mark.parametrize("activation_dtype", [torch.float32, torch.bfloat16]) +def test_kda_replay_reference_fold_and_graph( + dtype, capacity, width, projection_mode, lower_bound, gate_shift, activation_dtype +): + test_replay_acceptance_flush_and_graph( + width, + (128, 128), + capacity, + dtype, + run_config={"BV": 32, "num_warps": 2, "num_stages": 1}, + projection_mode=projection_mode, + gate_shift=gate_shift, + activation_dtype=activation_dtype, + head_geometry=(2, 2), + kda=True, + lower_bound=lower_bound, + ) + + +@pytest.mark.parametrize("dtype", [torch.float32, torch.bfloat16]) +def test_kda_compact_reference_and_partial_acceptance(dtype): + test_compact_preserves_accepted_prefix(dtype, kda=True) + + +@pytest.mark.parametrize("config", get_configs()) +@pytest.mark.parametrize("mode", ["compact", "inline", "precompute"]) +def test_kda_tuning_candidates_match_reference(config, mode): + if mode == "compact": + test_compact_preserves_accepted_prefix(torch.bfloat16, run_config=config, kda=True) + else: + test_replay_acceptance_flush_and_graph( + 3, + (128, 128), + 4, + torch.bfloat16, + run_config=config, + projection_mode=mode, + head_geometry=(2, 2), + kda=True, + gate_shift=-3, + ) diff --git a/unit_tests/common/basemodel/triton_kernel/linear_att/test_replayssm_execution.py b/unit_tests/common/basemodel/triton_kernel/linear_att/test_replayssm_execution.py index 99cf56f854..1b6c479978 100644 --- a/unit_tests/common/basemodel/triton_kernel/linear_att/test_replayssm_execution.py +++ b/unit_tests/common/basemodel/triton_kernel/linear_att/test_replayssm_execution.py @@ -13,12 +13,22 @@ @pytest.mark.parametrize("dtype", [torch.float32, torch.bfloat16]) @pytest.mark.parametrize("capture_metadata", [False, True]) @pytest.mark.parametrize("varlen", [False, True]) -def test_request_manager_accept_with_replay_graph(dtype, capture_metadata, varlen): +@pytest.mark.parametrize("kda", [False, True], ids=["gdn", "kda"]) +def test_request_manager_accept_with_replay_graph(dtype, capture_metadata, varlen, kda): torch.manual_seed(827) config = {"BV": 32, "num_warps": 1, "num_stages": 1} initial = torch.randn(3, 5, 4, 128, 128, device="cuda", dtype=dtype) * 0.01 caches = [ - cls(initial.clone(), 8, 4, torch.bfloat16, num_key_heads=2, projection_mode="precompute", run_config=config) + cls( + initial.clone(), + 8, + 4, + torch.bfloat16, + num_key_heads=2, + kda=kda, + projection_mode="precompute", + run_config=config, + ) for cls in [ReplaySSMCache, ReplaySSMCache] ] manager = object.__new__(ReqManagerForMamba) @@ -32,10 +42,10 @@ def test_request_manager_accept_with_replay_graph(dtype, capture_metadata, varle q = torch.randn(1, 12, 2, 128, device="cuda", dtype=torch.bfloat16) k = torch.randn_like(q) v = torch.randn(1, 12, 4, 128, device="cuda", dtype=torch.bfloat16) - a = torch.full((12, 4), -3.0, device="cuda", dtype=torch.bfloat16) - beta = torch.randn_like(a) + a = torch.full((12, 4 * 128 if kda else 4), -3.0, device="cuda", dtype=torch.bfloat16) + beta = torch.randn(a.shape[0], 4, device="cuda", dtype=a.dtype) alog = torch.zeros(4, device="cuda") - bias = torch.zeros_like(alog) + bias = torch.zeros(4 * 128 if kda else 4, device="cuda") def inputs(iteration): ids = [0, 1, 3] if iteration % 2 == 0 else [3, 0, 1] @@ -127,12 +137,19 @@ def forward_accept(index, cu, reqs): @pytest.mark.parametrize("dtype", [torch.float32, torch.bfloat16]) @pytest.mark.parametrize("capture_metadata", [False, True]) @pytest.mark.parametrize("stagger", [False, True]) -def test_decode_accept_order_preserves_forward_snapshot(dtype, capture_metadata, stagger): +@pytest.mark.parametrize("kda", [False, True], ids=["gdn", "kda"]) +def test_decode_accept_order_preserves_forward_snapshot(dtype, capture_metadata, stagger, kda): torch.manual_seed(731) initial = torch.randn(3, 3, 4, 128, 128, device="cuda", dtype=dtype) * 0.01 caches = [ ReplaySSMCache( - initial.clone(), 8, 1, dtype, num_key_heads=2, run_config={"BV": 32, "num_warps": 1, "num_stages": 1} + initial.clone(), + 8, + 1, + dtype, + num_key_heads=2, + kda=kda, + run_config={"BV": 32, "num_warps": 1, "num_stages": 1}, ) for _ in range(2) ] @@ -142,10 +159,10 @@ def test_decode_accept_order_preserves_forward_snapshot(dtype, capture_metadata, q = torch.randn(3, 2, 128, device="cuda", dtype=dtype) k = torch.randn_like(q) v = torch.randn(3, 4, 128, device="cuda", dtype=dtype) - a = torch.full((3, 4), -3.0, device="cuda", dtype=dtype) - beta = torch.randn_like(a) + a = torch.full((3, 4 * 128 if kda else 4), -3.0, device="cuda", dtype=dtype) + beta = torch.randn(a.shape[0], 4, device="cuda", dtype=a.dtype) alog = torch.zeros(4, device="cuda") - bias = torch.zeros_like(alog) + bias = torch.zeros(4 * 128 if kda else 4, device="cuda") def forward(cache): return [cache.forward(layer, q, k, v, a, beta, alog, bias, reqs) for layer in range(3)] @@ -257,12 +274,13 @@ def step(cache, accepted): @pytest.mark.parametrize("batch", [1, 8, 12]) @pytest.mark.parametrize("dtype", [torch.float32, torch.bfloat16]) -def test_small_batch_matches_larger_graph_padding(batch, dtype): +@pytest.mark.parametrize("kda", [False, True], ids=["gdn", "kda"]) +def test_small_batch_matches_larger_graph_padding(batch, dtype, kda): torch.manual_seed(1087 + batch) initial = torch.randn(2, 14, 4, 128, 128, device="cuda", dtype=dtype) * 0.01 config = {"BV": 32, "num_warps": 1, "num_stages": 1} caches = [ - ReplaySSMCache(initial.clone(), 8, 4, num_key_heads=2, projection_mode="precompute", run_config=config) + ReplaySSMCache(initial.clone(), 8, 4, num_key_heads=2, kda=kda, projection_mode="precompute", run_config=config) for _ in range(2) ] reqs = torch.arange(batch, device="cuda", dtype=torch.int32) @@ -272,10 +290,10 @@ def test_small_batch_matches_larger_graph_padding(batch, dtype): q = torch.randn(1, batch * 4, 2, 128, device="cuda", dtype=torch.bfloat16) k = torch.randn_like(q) v = torch.randn(1, batch * 4, 4, 128, device="cuda", dtype=torch.bfloat16) - a = torch.full((batch * 4, 4), -3.0, device="cuda", dtype=torch.bfloat16) - beta = torch.randn_like(a) + a = torch.full((batch * 4, 4 * 128 if kda else 4), -3.0, device="cuda", dtype=torch.bfloat16) + beta = torch.randn(a.shape[0], 4, device="cuda", dtype=a.dtype) alog = torch.zeros(4, device="cuda") - bias = torch.zeros_like(alog) + bias = torch.zeros(4 * 128 if kda else 4, device="cuda") accepted = torch.zeros(14, device="cuda", dtype=torch.int32) graphs, outputs = [], [] for cache, ids, offsets in zip(caches, [reqs, padded], [cu, padded_cu]): diff --git a/unit_tests/common/basemodel/triton_kernel/linear_att/test_ssm_autotune.py b/unit_tests/common/basemodel/triton_kernel/linear_att/test_ssm_autotune.py index 199bdc4ae9..fc7da18b36 100644 --- a/unit_tests/common/basemodel/triton_kernel/linear_att/test_ssm_autotune.py +++ b/unit_tests/common/basemodel/triton_kernel/linear_att/test_ssm_autotune.py @@ -13,20 +13,20 @@ from lightllm.common.triton_utils.autotuner import Autotuner, AutotuneKernelType -def make_inputs(mode, dtype=torch.bfloat16, device="cpu", width=4, projection_mode="inline"): +def make_inputs(mode, dtype=torch.bfloat16, device="cpu", width=4, projection_mode="inline", kda=False): layers, slots, batch, h, hv, kd, vd = 2, 9, 3, 2, 4, 32, 64 state = torch.randn(layers, slots, hv, kd, vd, dtype=dtype, device=device) * 0.01 if mode == "replay": - cache = ReplaySSMCache(state, 8, width, num_key_heads=h, projection_mode=projection_mode) + cache = ReplaySSMCache(state, 8, width, num_key_heads=h, projection_mode=projection_mode, kda=kda) else: - cache = CompactSSMCache(state, width, torch.bfloat16, num_key_heads=h) + cache = CompactSSMCache(state, width, torch.bfloat16, num_key_heads=h, kda=kda) tokens = batch * width packed = torch.randn(1, tokens, 2 * h * kd + hv * vd, dtype=torch.bfloat16, device=device) q, k, v = [ x.view(1, tokens, heads, dim) for x, heads, dim in zip(packed.split([h * kd, h * kd, hv * vd], -1), [h, h, hv], [kd, kd, vd]) ] - gates = torch.randn(tokens, hv * 2, device=device, dtype=q.dtype) + gates = torch.randn(tokens, hv * (kd + 1) if kda else hv * 2, device=device, dtype=q.dtype) a, b = gates[:, :-hv], gates[:, -hv:] return dict( cache=cache, @@ -77,8 +77,9 @@ def isolate_tuner(monkeypatch, tmp_path): @pytest.mark.parametrize("mode", ["gdn", "replay"]) @pytest.mark.parametrize("dtype", [torch.float32, torch.bfloat16]) -def test_rebuild_has_real_requests_and_independent_packed_inputs(mode, dtype): - inputs = make_inputs(mode, dtype) +@pytest.mark.parametrize("kda", [False, True], ids=["gdn", "kda"]) +def test_rebuild_has_real_requests_and_independent_packed_inputs(mode, dtype, kda): + inputs = make_inputs(mode, dtype, kda=kda) cache = inputs["cache"] for x in tensors(cache).values(): x.fill_(1) @@ -101,9 +102,9 @@ def test_rebuild_has_real_requests_and_independent_packed_inputs(mode, dtype): assert torch.all(x == 1) key = tuning.select_config._static_key(**inputs) assert len(KernelConfigs.get_config_file_name(key).encode()) <= 255 - other = make_inputs(mode, torch.float32 if dtype == torch.bfloat16 else torch.bfloat16) + other = make_inputs(mode, torch.float32 if dtype == torch.bfloat16 else torch.bfloat16, kda=kda) assert tuning.select_config._static_key(**other) != key - other = make_inputs(mode, dtype, width=3) + other = make_inputs(mode, dtype, width=3, kda=kda) assert tuning.select_config._static_key(**other) != key inputs["q"] = inputs["q"].contiguous() assert tuning.select_config._static_key(**inputs) != key @@ -145,8 +146,9 @@ def test_layout_is_fixed_before_graph_capture_and_cached_configs_do_not_execute( ], ) @pytest.mark.parametrize("dtype", [torch.float32, torch.bfloat16]) -def test_joint_tuning_isolated_state_and_graph_replay(monkeypatch, mode, dtype, width, projection_mode): - inputs = make_inputs(mode, dtype, "cuda", width=width, projection_mode=projection_mode) +@pytest.mark.parametrize("kda", [False, True], ids=["gdn", "kda"]) +def test_joint_tuning_isolated_state_and_graph_replay(monkeypatch, mode, dtype, width, projection_mode, kda): + inputs = make_inputs(mode, dtype, "cuda", width=width, projection_mode=projection_mode, kda=kda) if width == 1: for name in ("q", "k", "v"): inputs[name] = inputs[name].transpose(0, 1) @@ -182,7 +184,7 @@ def checked_bench(*args, **kwargs): assert cache.run_config in configs # An explicit identical layout supplies a reference for capture/replay and # alternating real/HOLD rows. It must never trigger its own tuning. - expected = make_inputs(mode, dtype, "cuda", width=width, projection_mode=projection_mode)["cache"] + expected = make_inputs(mode, dtype, "cuda", width=width, projection_mode=projection_mode, kda=kda)["cache"] expected.run_config = cache.run_config.copy() expected._config_is_fixed = True expected.state.copy_(before["state"]) @@ -219,3 +221,16 @@ def step(target): torch.testing.assert_close(cache.state, expected.state, atol=0, rtol=0) torch.testing.assert_close(cache.state[:, cache.hold], before["state"][:, cache.hold], atol=0, rtol=0) assert len(timings) == len(configs) + + +@pytest.mark.parametrize("mode", ["gdn", "replay"]) +def test_kda_rule_and_gate_bound_separate_tuning_cache(mode): + inputs = make_inputs(mode, kda=True) + key = tuning.select_config._static_key(**inputs) + assert key != tuning.select_config._static_key(**make_inputs(mode)) + _, rebuilt = tuning.rebuild_inputs(**inputs) + assert rebuilt["cache"].kda + assert rebuilt["cache"].lower_bound == inputs["cache"].lower_bound + assert key == tuning.select_config._static_key(**rebuilt) + inputs["cache"].lower_bound = -2.5 + assert key != tuning.select_config._static_key(**inputs) diff --git a/unit_tests/models/glm5_next/test_cache.py b/unit_tests/models/glm5_next/test_cache.py index ae40702d75..ce461914f6 100644 --- a/unit_tests/models/glm5_next/test_cache.py +++ b/unit_tests/models/glm5_next/test_cache.py @@ -38,12 +38,14 @@ def test_nope_cache_config_uses_native_mla_width(): @pytest.mark.parametrize("small_page", [False, True]) @pytest.mark.parametrize("tp_world_size", [1, 4]) @pytest.mark.parametrize("mtp_step", [0, 2]) -def test_hybrid_checkpoint_restore_and_packed_kv_copy(monkeypatch, small_page, tp_world_size, mtp_step): +@pytest.mark.parametrize("mode", ["native", "compact", "replay"]) +def test_hybrid_checkpoint_restore_and_packed_kv_copy(monkeypatch, small_page, tp_world_size, mtp_step, mode): monkeypatch.setenv("LIGHTLLM_CURRENT_RANK_IN_NODE", "0") monkeypatch.setenv("LIGHTLLM_CURRENT_DEVICE_ID", "0") monkeypatch.setattr("lightllm.common.req_manager.req_sampling_params.get_vocab_size", lambda _: 128) args = StartArgs( tp=tp_world_size, + ssm_state_mode=mode, data_type="bfloat16", linear_att_hash_page_size=4, linear_att_page_block_num=2, @@ -99,7 +101,28 @@ def test_hybrid_checkpoint_restore_and_packed_kv_copy(monkeypatch, small_page, t req.req_to_conv_state.buffer[:, 0].normal_() req.req_to_ssm_state.buffer[:, 0].normal_() conv = req.req_to_conv_state.buffer[:, 0, ..., :3].clone() - ssm = req.req_to_ssm_state.buffer[:, 0].clone() + updates = req.ssm_update_cache + if updates is not None: + # Save an accepted suffix, not merely an already folded checkpoint. + heads, dim = config.num_linear_v_heads, config.head_linear_k_dim + width = mtp_step + 1 + q = torch.randn(1, width, heads, dim, dtype=torch.bfloat16, device="cuda") + gate = torch.randn(width, heads * dim, dtype=torch.bfloat16, device="cuda") + beta = torch.randn(width, heads, dtype=torch.bfloat16, device="cuda") + alog = torch.zeros(heads, device="cuda") + bias = torch.zeros(heads * dim, device="cuda") + ids = torch.tensor([0], dtype=torch.int32, device="cuda") + cu = torch.tensor([0, width], dtype=torch.int32, device="cuda") + updates.prepare_decode(ids, cu) + for layer in range(config.linear_layer_num): + updates.forward(layer, q, q, q, gate, beta, alog, bias, ids, cu) + accepted = torch.zeros(4, dtype=torch.int32, device="cuda") + updates.accept_updates(ids, accepted) + ssm = updates.snapshot_accepted_state(0).clone() + if mtp_step: + req.req_to_mtp_state_index[2] = mtp_step + else: + ssm = req.req_to_ssm_state.buffer[:, 0].clone() if small_page: req.save_state(0, slot, cache) else: @@ -115,8 +138,14 @@ def test_hybrid_checkpoint_restore_and_packed_kv_copy(monkeypatch, small_page, t else: req.restore_big_page_state(slot, dest_req) torch.cuda.synchronize() + if updates is not None: + assert updates.kda + if mode == "replay": + assert updates.cursors[2] == 0 + if mtp_step: + assert req.req_to_mtp_state_index[2] == 0 assert torch.equal(req.req_to_conv_state.buffer[:, 2, ..., :3], conv) - assert torch.equal(req.req_to_ssm_state.buffer[:, 2 * (mtp_step + 1)], ssm) + assert torch.equal(req.req_to_ssm_state.buffer[:, 2 * req.ssm_slots_per_req], ssm) # Both page sizes are aligned to complete pools. Restoring a prefix must # discard stale tail values from a previously allocated request slot. assert not req.req_to_indexer_tail.buffer[:, 2].any() @@ -124,7 +153,7 @@ def test_hybrid_checkpoint_restore_and_packed_kv_copy(monkeypatch, small_page, t req.req_to_indexer_tail.buffer[:, 2].normal_() req.init_hybrid_attention_state(dest_req) assert not req.req_to_conv_state.buffer[:, 2].any() - assert not req.req_to_ssm_state.buffer[:, 2 * (mtp_step + 1) : 3 * (mtp_step + 1)].any() + assert not req.req_to_ssm_state.buffer[:, 2 * req.ssm_slots_per_req : 3 * req.ssm_slots_per_req].any() assert not req.req_to_indexer_tail.buffer[:, 2].any() assert torch.equal(req.req_to_indexer_tail.buffer[:, [0, 1, 3]], unchanged) # KV moves carry MLA latents and pooled FP8 bytes; raw keys/gates only diff --git a/unit_tests/models/glm5_next/test_mtp.py b/unit_tests/models/glm5_next/test_mtp.py index dac19a44b5..0fe7514fc7 100644 --- a/unit_tests/models/glm5_next/test_mtp.py +++ b/unit_tests/models/glm5_next/test_mtp.py @@ -122,7 +122,8 @@ def prefill(values, destination, tail): @pytest.mark.parametrize("dynamic", [False, True]) -def test_kda_backend_routes_accepted_conv_and_ssm_states(dynamic): +@pytest.mark.parametrize("mode,mtp_step", [("native", 2), ("compact", 2), ("replay", 2), ("native", 0), ("replay", 0)]) +def test_kda_backend_routes_accepted_conv_and_ssm_states(dynamic, mode, mtp_step): from types import SimpleNamespace from lightllm.common.basemodel.attention.base_att import AttControl @@ -130,20 +131,20 @@ def test_kda_backend_routes_accepted_conv_and_ssm_states(dynamic): from lightllm.common.basemodel.triton_kernel.linear_att.causal_conv1d import causal_conv1d_update torch.manual_seed(53) - heads, dim, width = 2, 128, 3 + heads, dim, width = 2, 128, mtp_step + 1 hidden = heads * dim - lengths = [2, 1, 3] if dynamic else [3, 3, 3] + lengths = ([2, 1, 3] if dynamic else [3, 3, 3]) if mtp_step else [1, 1, 1] requests = [2, 0, 1] request_rows = [req for req, length in zip(requests, lengths) for _ in range(length)] offsets = [offset for length in lengths for offset in range(length)] - if dynamic: + if dynamic and mtp_step: request_rows += [3, 3] offsets += [0, 0] tokens = len(request_rows) - conv = torch.randn(4, 3 * hidden, 5, device="cuda", dtype=torch.bfloat16) + conv = torch.randn(4, 3 * hidden, 3 + mtp_step, device="cuda", dtype=torch.bfloat16) ssm = torch.randn(4 * width, heads, dim, dim, device="cuda") * 0.01 old_conv, expected_ssm = conv.clone(), ssm.clone() - accepted_offsets = torch.tensor([1, 2, 0, 0], dtype=torch.int32, device="cuda") + accepted_offsets = torch.tensor([1, 2, 0, 0] if mtp_step else [0, 0, 0, 0], dtype=torch.int32, device="cuda") mixed = torch.randn(tokens, 3 * hidden, device="cuda", dtype=torch.bfloat16) gate = torch.randn(tokens, hidden, device="cuda", dtype=torch.bfloat16) beta = torch.randn(tokens, heads, device="cuda", dtype=torch.bfloat16) @@ -181,7 +182,7 @@ def test_kda_backend_routes_accepted_conv_and_ssm_states(dynamic): expected_ssm[req * width + token] = state[0] row += 1 backend = SimpleNamespace( - mtp_step=2, + mtp_step=mtp_step, tp_num_heads=heads, head_dim=dim, tp_hidden_size=hidden, @@ -190,8 +191,25 @@ def test_kda_backend_routes_accepted_conv_and_ssm_states(dynamic): split_qkv=lambda x: x.split(hidden, -1), ) manager = SimpleNamespace( - req_to_mtp_state_index=accepted_offsets, HOLD_REQUEST_ID=3, get_mamba_cache=lambda layer: (conv, ssm) + req_to_mtp_state_index=accepted_offsets, + HOLD_REQUEST_ID=3, + get_mamba_cache=lambda layer: (conv, ssm), + ssm_update_cache=None, + ssm_slots_per_req=width, + linear_config=SimpleNamespace(full_attention_interval=4), ) + if mode != "native": + from lightllm.common.basemodel.triton_kernel.linear_att.replayssm import ReplaySSMCache + from lightllm.common.basemodel.triton_kernel.linear_att.replayssm_compact import CompactSSMCache + + indexes = torch.arange(4, device="cuda") * width + accepted_offsets + ssm = ssm[indexes].clone() + before_verify = ssm.clone() + manager.ssm_slots_per_req = 1 + if mode == "replay": + manager.ssm_update_cache = ReplaySSMCache(ssm.unsqueeze(0), 8, width, kda=True) + else: + manager.ssm_update_cache = CompactSSMCache(ssm.unsqueeze(0), width, torch.bfloat16, kda=True) infer = SimpleNamespace( batch_size=tokens, b_req_idx=torch.tensor(request_rows, device="cuda", dtype=torch.int32), @@ -217,6 +235,107 @@ def test_kda_backend_routes_accepted_conv_and_ssm_states(dynamic): ), ) torch.testing.assert_close(actual.reshape(tokens, heads, dim)[:row], torch.cat(expected), atol=2e-3, rtol=1e-2) - torch.testing.assert_close(ssm, expected_ssm, atol=2e-3, rtol=1e-2) - if dynamic: + if mode == "native": + torch.testing.assert_close(ssm, expected_ssm, atol=2e-3, rtol=1e-2) + else: + assert torch.equal(ssm, before_verify) + accepted = torch.zeros(4, device="cuda", dtype=torch.int32) + for req, length in zip(requests, lengths): + accepted[req] = min(length, 2) - 1 + if mtp_step: + manager.ssm_update_cache.accept_updates(state.b_conv_buffer_idx, accepted) + manager.ssm_update_cache.merge_accepted_updates(state.b_conv_buffer_idx) + for req in requests: + torch.testing.assert_close(ssm[req], expected_ssm[req * width + accepted[req]], atol=2e-3, rtol=1e-2) + if dynamic and mtp_step: assert torch.equal(conv[3], old_conv[3]) + + +@pytest.mark.parametrize("mode", ["compact", "replay"]) +def test_kda_prefill_materializes_accepted_history_and_uses_single_state_slot(mode): + import triton + from types import SimpleNamespace + from lightllm.common.basemodel.attention.base_att import AttControl + from lightllm.common.basemodel.attention.linear.kda import KDAPrefillAttState + from lightllm.common.basemodel.triton_kernel.linear_att.replayssm import ReplaySSMCache + from lightllm.common.basemodel.triton_kernel.linear_att.replayssm_compact import CompactSSMCache + + torch.manual_seed(1353) + triton.set_allocator(lambda size, alignment, stream: torch.empty(size, device="cuda", dtype=torch.int8)) + heads, dim, width, tokens = 2, 128, 3, 5 + hidden = heads * dim + checkpoint = torch.randn(1, 4, heads, dim, dim, device="cuda") * 0.01 + if mode == "replay": + updates = ReplaySSMCache(checkpoint, 8, width, kda=True) + else: + updates = CompactSSMCache(checkpoint, width, torch.bfloat16, kda=True) + ids = torch.tensor([1], dtype=torch.int32, device="cuda") + verify = torch.randn(1, width, heads, dim, dtype=torch.bfloat16, device="cuda") + gate = torch.randn(width, hidden, dtype=torch.bfloat16, device="cuda") + beta = torch.randn(width, heads, dtype=torch.bfloat16, device="cuda") + alog = torch.randn(heads, device="cuda") * 0.1 + bias = torch.randn(hidden, device="cuda") * 0.1 + updates.prepare_decode(ids, ids.new_tensor([0, width])) + updates.forward(0, verify, verify, verify, gate, beta, alog, bias, ids, ids.new_tensor([0, width])) + updates.accept_updates(ids, ids.new_tensor([0, 1, 0, 0])) + canonical = updates.snapshot_accepted_state(1).clone() + if mode == "replay": + assert updates.cursors[1] == 2 + + conv = torch.randn(4, 3 * hidden, 5, dtype=torch.bfloat16, device="cuda") * 0.1 + native_conv = conv.clone() + native_ssm = torch.randn(4 * width, heads, dim, dim, device="cuda") + native_ssm[width].copy_(canonical[0]) + weight = SimpleNamespace( + get_merged_kda_conv_weight=lambda: conv_weight, + linear_A_log=SimpleNamespace(weight=alog), + linear_dt_bias=SimpleNamespace(weight=bias), + ) + conv_weight = torch.randn(3 * hidden, 4, dtype=torch.bfloat16, device="cuda") * 0.1 + backend = SimpleNamespace( + mtp_step=2, + tp_num_heads=heads, + head_dim=dim, + tp_hidden_size=hidden, + conv_kernel_size=4, + lower_bound=-5.0, + split_qkv=lambda x: x.split(hidden, -1), + ) + mixed = torch.randn(tokens, 3 * hidden, dtype=torch.bfloat16, device="cuda") + gate = torch.randn(tokens, hidden, dtype=torch.bfloat16, device="cuda") + beta = torch.randn(tokens, heads, dtype=torch.bfloat16, device="cuda") + outputs = [] + for cache, conv_pool, ssm_pool, slots in [ + (None, native_conv, native_ssm, width), + (updates, conv, checkpoint[0], 1), + ]: + manager = SimpleNamespace( + ssm_update_cache=cache, ssm_slots_per_req=slots, get_mamba_cache=lambda layer: (conv_pool, ssm_pool) + ) + infer = SimpleNamespace( + b_req_idx=ids, + b_ready_cache_len=ids.new_tensor([2]), + b1_cu_q_seq_len=ids.new_tensor([0, tokens]), + req_manager=manager, + ) + state = KDAPrefillAttState(backend=backend, infer_state=infer) + state.init_state() + assert state.b_ssm_buffer_idx.tolist() == [slots] + outputs.append( + state.prefill_att( + None, + None, + None, + AttControl( + linear_att_prefill=True, + linear_att_prefill_dict=dict( + layer_weight=weight, layer_num=0, mixed_qkv=mixed.clone(), raw_gate=gate, raw_beta=beta + ), + ), + ) + ) + torch.testing.assert_close(outputs[1], outputs[0], atol=0, rtol=0) + torch.testing.assert_close(checkpoint[0, 1], native_ssm[width], atol=0, rtol=0) + torch.testing.assert_close(conv[1], native_conv[1], atol=0, rtol=0) + if mode == "replay": + assert updates.cursors[1] == 0 diff --git a/unit_tests/models/glm5_next/test_pd_cache.py b/unit_tests/models/glm5_next/test_pd_cache.py index 57665cbbe3..4983679b69 100644 --- a/unit_tests/models/glm5_next/test_pd_cache.py +++ b/unit_tests/models/glm5_next/test_pd_cache.py @@ -14,12 +14,12 @@ pytestmark = pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") -@pytest.fixture(params=[0, 2]) +@pytest.fixture(params=[(0, "native"), (2, "native"), (0, "replay"), (2, "replay"), (2, "compact")]) def make_mems(monkeypatch, request): monkeypatch.setenv("LIGHTLLM_CURRENT_RANK_IN_NODE", "0") monkeypatch.setenv("LIGHTLLM_CURRENT_DEVICE_ID", "0") monkeypatch.setattr("lightllm.common.req_manager.req_sampling_params.get_vocab_size", lambda _: 128) - mtp_step = request.param + mtp_step, state_mode = request.param args = StartArgs( data_type="bfloat16", linear_att_hash_page_size=4, @@ -39,7 +39,10 @@ def check_shm_refs(mem, req): monkeypatch.setattr(MemoryManager, "write_to_shm", check_shm_refs) - def create(tp): + def create(tp, mode=None): + args.ssm_state_mode = mode or state_mode + set_env_start_args(dataclasses.asdict(args)) + get_env_start_args.cache_clear() config = Glm5NextCacheConfig( tp_world_size=tp, full_att_all_num_kv_heads=1, @@ -64,6 +67,8 @@ def create(tp): for _ in range(tp): mem = Glm5NextMemManager(32, torch.bfloat16, 1, 584, 2 + int(mtp_step > 0), config) req = Glm5NextReqManager(3, 32, mem, config) + assert req.ssm_slots_per_req == (mtp_step + 1 if args.ssm_state_mode == "native" else 1) + assert (req.ssm_update_cache is None) == (args.ssm_state_mode == "native") mem.write_to_shm(req) assert mem.big_page_buffers is not None mems.append(mem) @@ -75,8 +80,9 @@ def create(tp): @pytest.mark.parametrize("prefill_tp,decode_tp", [(1, 1), (4, 4), (1, 4), (4, 1)]) @pytest.mark.parametrize("remainder", range(4)) -def test_pd_kv_and_runtime_state_roundtrip(make_mems, prefill_tp, decode_tp, remainder): - source, dest = make_mems(prefill_tp), make_mems(decode_tp) +@pytest.mark.parametrize("source_mode", ["native", "replay"]) +def test_pd_kv_and_runtime_state_roundtrip(make_mems, prefill_tp, decode_tp, remainder, source_mode): + source, dest = make_mems(prefill_tp, source_mode), make_mems(decode_tp) mtp_size = get_env_start_args().mtp_step + 1 layers = source[0].layer_num src_req, dst_req = 0, 2 @@ -93,8 +99,33 @@ def test_pd_kv_and_runtime_state_roundtrip(make_mems, prefill_tp, decode_tp, rem heads = slice(rank * 8 // prefill_tp, (rank + 1) * 8 // prefill_tp) mem.kv_buffer.view(torch.uint8)[:, src_indexes] = packed mem.req_to_conv_state.buffer[:, src_req, ..., :3].copy_(conv[:, :, heads].reshape(6, -1, 3)) - mem.req_to_ssm_state.buffer[:, src_req * mtp_size].copy_(ssm[:, heads]) + mem.req_to_ssm_state.buffer[:, src_req * mem.ssm_slots_per_req].copy_(ssm[:, heads]) mem.req_to_indexer_tail.buffer[:, src_req].copy_(tail) + updates = mem.ssm_update_cache + if updates is not None: + # The exported checkpoint must contain accepted, unfurled KDA history. + h = mem.linear_config.num_linear_v_heads + q = torch.randn(1, 1, h, 128, dtype=torch.bfloat16, device="cuda") + gate = torch.randn(1, h * 128, dtype=torch.bfloat16, device="cuda") + beta = torch.randn(1, h, dtype=torch.bfloat16, device="cuda") + reqs = torch.tensor([src_req], dtype=torch.int32, device="cuda") + cu = torch.tensor([0, 1], dtype=torch.int32, device="cuda") + updates.prepare_decode(reqs, cu) + for layer in range(6): + updates.forward( + layer, + q, + q, + q, + gate, + beta, + torch.zeros(h, device="cuda"), + torch.zeros(h * 128, device="cuda"), + reqs, + cu, + ) + updates.accept_updates(reqs, torch.zeros(4, dtype=torch.int32, device="cuda")) + ssm[:, heads].copy_(updates.snapshot_accepted_state(src_req)) untouched = [] for mem in dest: @@ -102,12 +133,16 @@ def test_pd_kv_and_runtime_state_roundtrip(make_mems, prefill_tp, decode_tp, rem mem.req_to_conv_state.buffer.normal_() mem.req_to_ssm_state.buffer.normal_() mem.req_to_indexer_tail.buffer.normal_() + if mem.ssm_update_cache is not None and hasattr(mem.ssm_update_cache, "cursors"): + mem.ssm_update_cache.cursors[dst_req] = 5 + if mem.req_to_mtp_state_index is not None: + mem.req_to_mtp_state_index[dst_req] = mtp_size - 1 untouched.append( [ x.buffer[:, torch.tensor([0, 1, 3], device="cuda") * stride].clone() for x, stride in ( (mem.req_to_conv_state, 1), - (mem.req_to_ssm_state, mtp_size), + (mem.req_to_ssm_state, mem.ssm_slots_per_req), (mem.req_to_indexer_tail, 1), ) ] @@ -126,10 +161,14 @@ def test_pd_kv_and_runtime_state_roundtrip(make_mems, prefill_tp, decode_tp, rem heads = slice(rank * 8 // decode_tp, (rank + 1) * 8 // decode_tp) assert torch.equal(mem.kv_buffer.view(torch.uint8)[:, dst_indexes], packed) assert torch.equal(mem.req_to_conv_state.buffer[:, dst_req, ..., :3], conv[:, :, heads].reshape(6, -1, 3)) - assert torch.equal(mem.req_to_ssm_state.buffer[:, dst_req * mtp_size], ssm[:, heads]) + assert torch.equal(mem.req_to_ssm_state.buffer[:, dst_req * mem.ssm_slots_per_req], ssm[:, heads]) assert torch.equal(mem.req_to_indexer_tail.buffer[:, dst_req], tail) + if mem.ssm_update_cache is not None and hasattr(mem.ssm_update_cache, "cursors"): + assert mem.ssm_update_cache.cursors[dst_req] == 0 + if mem.req_to_mtp_state_index is not None and mem.ssm_update_cache is not None: + assert mem.req_to_mtp_state_index[dst_req] == 0 for (state, stride), expected in zip( - ((mem.req_to_conv_state, 1), (mem.req_to_ssm_state, mtp_size), (mem.req_to_indexer_tail, 1)), + ((mem.req_to_conv_state, 1), (mem.req_to_ssm_state, mem.ssm_slots_per_req), (mem.req_to_indexer_tail, 1)), untouched[rank], ): indexes = torch.tensor([0, 1, 3], device="cuda") * stride diff --git a/unit_tests/server/test_ssm_state_mode.py b/unit_tests/server/test_ssm_state_mode.py index 672c063515..b716850207 100644 --- a/unit_tests/server/test_ssm_state_mode.py +++ b/unit_tests/server/test_ssm_state_mode.py @@ -12,6 +12,9 @@ ("compact", "bfloat16", 2, "qwen3_5", 16, None), ("compact", "float32", 2, "qwen3_5", 16, None), ("compact", "float32", 32, "qwen3_5", 16, None), + ("compact", "bfloat16", 2, "glm5_next", 16, None), + ("replay", "float32", 0, "glm5_next", 8, None), + ("replay", "bfloat16", 2, "glm5_next_text", 8, None), ("replay", "float32", 0, "qwen3_next", 16, None), ("replay", "float32", 15, "qwen3_5", 16, None), ("compact", "bfloat16", 0, "qwen3_5", 16, "requires MTP"), From 28f2acaf673d80bf55fc44d2ffe880a85f4361bc Mon Sep 17 00:00:00 2001 From: sufubao Date: Tue, 6 Oct 2026 13:56:01 +0800 Subject: [PATCH 06/40] fix: keep KDA tuning cache filenames within filesystem limits --- docs/replayssm/README.md | 2 +- .../triton_kernel/linear_att/ssm_autotune.py | 2 +- .../linear_att/test_ssm_autotune.py | 27 +++++++++++++++++++ 3 files changed, 29 insertions(+), 2 deletions(-) diff --git a/docs/replayssm/README.md b/docs/replayssm/README.md index 1e9dd560a0..f80c6333c1 100644 --- a/docs/replayssm/README.md +++ b/docs/replayssm/README.md @@ -7,7 +7,7 @@ the default. ## Configuration ```bash -python -m lightllm.server.api_server \ +LIGHTLLM_TRITON_AUTOTUNE_LEVEL=1 python -m lightllm.server.api_server \ --model_dir /path/to/qwen3.5 \ --ssm_state_mode replay \ --linear_att_ssm_data_type bfloat16 \ diff --git a/lightllm/common/basemodel/triton_kernel/linear_att/ssm_autotune.py b/lightllm/common/basemodel/triton_kernel/linear_att/ssm_autotune.py index d1b3fac568..032c7daeb7 100644 --- a/lightllm/common/basemodel/triton_kernel/linear_att/ssm_autotune.py +++ b/lightllm/common/basemodel/triton_kernel/linear_att/ssm_autotune.py @@ -37,7 +37,7 @@ def static_key(cache, mode, q, k, v, a, b, cu_seqlens): key["history_dtype"] = str(cache.keys.dtype) key["v"] = 4 # Small batches omit the forward fold already handled by prepare. key["projection"] = ( - "block_mma_activation_precision" + ("kda_mma" if cache.kda else "block_mma_activation_precision") if cache.projection_mode == "precompute" and cache.verify_width > 1 else "inline" ) diff --git a/unit_tests/common/basemodel/triton_kernel/linear_att/test_ssm_autotune.py b/unit_tests/common/basemodel/triton_kernel/linear_att/test_ssm_autotune.py index fc7da18b36..396c3e863d 100644 --- a/unit_tests/common/basemodel/triton_kernel/linear_att/test_ssm_autotune.py +++ b/unit_tests/common/basemodel/triton_kernel/linear_att/test_ssm_autotune.py @@ -234,3 +234,30 @@ def test_kda_rule_and_gate_bound_separate_tuning_cache(mode): assert key == tuning.select_config._static_key(**rebuilt) inputs["cache"].lower_bound = -2.5 assert key != tuning.select_config._static_key(**inputs) + + +@pytest.mark.parametrize("tp", [1, 2, 4, 8]) +@pytest.mark.parametrize("dtype", [torch.float32, torch.bfloat16]) +def test_kda_production_tuning_cache_can_be_written(monkeypatch, tmp_path, tp, dtype): + from types import SimpleNamespace + + monkeypatch.setattr(kernel_config_module, "get_current_device_name", lambda: "NVIDIA H100 80GB HBM3") + heads = 64 // tp + cache = SimpleNamespace( + state=SimpleNamespace(shape=(34, 1, heads, 128, 128), dtype=dtype), + keys=SimpleNamespace(dtype=torch.bfloat16), + verify_width=3, + num_key_heads=heads, + capacity=8, + kda=True, + lower_bound=-5.0, + projection_mode="precompute", + ) + packed = torch.empty(1, 3, heads * 128 * 3, dtype=torch.bfloat16) + q, k, v = [x.view(1, 3, heads, 128) for x in packed.chunk(3, -1)] + a = torch.empty(3, heads * 128, dtype=q.dtype) + b = torch.empty(3, heads, dtype=q.dtype) + key = tuning.static_key(cache, "replay", q, k, v, a, b, torch.tensor([0, 3])) + path = tmp_path / KernelConfigs.get_config_file_name(key) + path.write_text("{}") + assert path.is_file() From 59fae105572ee6b585800063544e312b29b109e0 Mon Sep 17 00:00:00 2001 From: sufubao Date: Tue, 6 Oct 2026 14:50:35 +0800 Subject: [PATCH 07/40] fix: handle empty GLM indexer rows during graph capture --- lightllm/models/glm5_next/indexer.py | 7 +++++-- .../models/glm5_next/triton_kernel/kpool.py | 4 ++-- .../models/glm5_next/test_indexer_topk.py | 17 +++++++++-------- unit_tests/models/glm5_next/test_kernels.py | 9 ++++++--- 4 files changed, 22 insertions(+), 15 deletions(-) diff --git a/lightllm/models/glm5_next/indexer.py b/lightllm/models/glm5_next/indexer.py index cf13ab83de..2acc51ca94 100644 --- a/lightllm/models/glm5_next/indexer.py +++ b/lightllm/models/glm5_next/indexer.py @@ -113,12 +113,15 @@ def _get_decode_indices(self, q_fp8, weights, packed_buffer, infer_state, att_st lengths, max_pools, ) - metadata = deep_gemm.get_paged_mqa_logits_metadata(lengths, 64, deep_gemm.get_num_sms()) + # DeepGEMM's SM90 scheduler cannot handle a large batch of empty HOLD rows. + # Their initialized dummy pool is discarded by top-k's original lengths. + schedule_lengths = lengths.clamp_min(1) + metadata = deep_gemm.get_paged_mqa_logits_metadata(schedule_lengths, 64, deep_gemm.get_num_sms()) # Each MTP position has its own pool length; HOLD rows have no valid pools. # Fixed-width mtp_step=2 gives Q=[3 * num_requests, 1, heads, dim], so next_n stays 1. # All three verify positions are preserved despite the SM90 kernel's native next_n limit of 2. logits = deep_gemm.fp8_paged_mqa_logits( - q_fp8.unsqueeze(1), pages, weights, lengths, block_table, metadata, max_pools, clean_logits=False + q_fp8.unsqueeze(1), pages, weights, schedule_lengths, block_table, metadata, max_pools, clean_logits=False ) groups = torch.empty((q_fp8.shape[0], self.topk // 4), dtype=torch.int32, device=q_fp8.device) self.select_topk_indices(logits, lengths.view(-1), groups) diff --git a/lightllm/models/glm5_next/triton_kernel/kpool.py b/lightllm/models/glm5_next/triton_kernel/kpool.py index 0e126aea6a..277189025a 100644 --- a/lightllm/models/glm5_next/triton_kernel/kpool.py +++ b/lightllm/models/glm5_next/triton_kernel/kpool.py @@ -276,7 +276,7 @@ def _gather_paged_pools( rows = tl.arange(0, 64) cols = tl.arange(0, 128) # A fixed grid is replayable at 1M; the GPU length bounds the work. - for page in range(tl.program_id(0), tl.cdiv(length, 64), tl.num_programs(0)): + for page in range(tl.program_id(0), tl.cdiv(tl.maximum(length, 1), 64), tl.num_programs(0)): pools = page * 64 + rows valid = pools < length locs = tl.load(ReqTable + req * REQ_STRIDE + pools * 4 + 3, valid, 0).to(tl.int64) @@ -289,7 +289,7 @@ def _gather_paged_pools( def gather_paged_pools(packed_buffer, req_table, req_idx, pool_lengths, max_pools): - """Pack valid pools into DeepGEMM pages; unused pages remain unread.""" + """Pack valid pools, including one zero-key dummy page for empty rows.""" batch = req_idx.numel() max_pages = triton.cdiv(max_pools, 64) pages = torch.empty((batch * max_pages, 64, 1, 132), device=packed_buffer.device, dtype=torch.uint8) diff --git a/unit_tests/models/glm5_next/test_indexer_topk.py b/unit_tests/models/glm5_next/test_indexer_topk.py index fce877448c..39bb6220ad 100644 --- a/unit_tests/models/glm5_next/test_indexer_topk.py +++ b/unit_tests/models/glm5_next/test_indexer_topk.py @@ -98,20 +98,19 @@ def test_topk_without_vllm(indexer, monkeypatch): @pytest.mark.parametrize("mtp_size", [1, 3, 6]) -@pytest.mark.parametrize("max_pools", [1024, 262144]) -def test_paged_decode_matches_nonpaged_with_graph_and_mtp(indexer, monkeypatch, mtp_size, max_pools): +@pytest.mark.parametrize("requests,max_pools", [(5, 640), (5, 1024), (5, 262144), (64, 640)]) +def test_paged_decode_matches_nonpaged_with_graph_and_mtp(indexer, monkeypatch, mtp_size, requests, max_pools): pytest.importorskip("deep_gemm") - _require_vllm_topk() torch.manual_seed(27) - rows = 5 * mtp_size + rows = requests * mtp_size storage = torch.empty(1024, 1, 584, device="cuda", dtype=torch.bfloat16) packed = storage.view(torch.uint8)[:, :, -132:] packed[:, 0, :128] = torch.randn(1024, 128, device="cuda").to(torch.float8_e4m3fn).view(torch.uint8) scales = torch.pow(2.0, torch.randint(-3, 2, (1024,), device="cuda").float()) packed[:, 0, 128:] = scales.view(torch.uint8).view(-1, 4) - table = torch.full((5, max_pools * 4), -1, device="cuda", dtype=torch.int32) - table[:, 3::4] = torch.randint(0, 1024, (5, max_pools), device="cuda", dtype=torch.int32) - req_idx = torch.tensor([1, 3, 0, 2, 4], device="cuda", dtype=torch.int32).repeat_interleave(mtp_size) + table = torch.full((requests, max_pools * 4), -1, device="cuda", dtype=torch.int32) + table[:, 3::4] = torch.randint(0, 1024, (requests, max_pools), device="cuda", dtype=torch.int32) + req_idx = torch.arange(requests, device="cuda", dtype=torch.int32).roll(2).repeat_interleave(mtp_size) lengths = torch.zeros(rows, device="cuda", dtype=torch.int32) q = torch.randn(rows, 32, 128, device="cuda").to(torch.float8_e4m3fn) weights = torch.randn(rows, 32, device="cuda") * 0.05 @@ -142,8 +141,10 @@ def run(): with torch.cuda.graph(graph): actual = run() graph_logits = logits_outputs.pop() + graph.replay() + assert (actual == -1).all() for bases in ([0, 2047, 2051, max_pools * 4 - mtp_size + 1, 0], [7, 65, 511, 2053, 0]): - seqs = torch.tensor(bases, device="cuda", dtype=torch.int32)[:, None] + seqs = torch.tensor(bases, device="cuda", dtype=torch.int32).repeat((requests + 4) // 5)[:requests, None] seqs = seqs + torch.arange(mtp_size, device="cuda", dtype=torch.int32)[None, :] seqs[-1].zero_() lengths.copy_(seqs.flatten()) diff --git a/unit_tests/models/glm5_next/test_kernels.py b/unit_tests/models/glm5_next/test_kernels.py index 3270ba972e..a792d1c080 100644 --- a/unit_tests/models/glm5_next/test_kernels.py +++ b/unit_tests/models/glm5_next/test_kernels.py @@ -322,7 +322,7 @@ def test_gather_paged_pools_valid_pages_and_graph_replay(max_pools): def check(pages, block_table): page_bytes = pages.view(pages.shape[0], -1) for row, (req, length) in enumerate(zip(req_idx.tolist(), lengths.tolist())): - page_ids = block_table[row, : triton.cdiv(length, 64)].long() + page_ids = block_table[row, : max(1, triton.cdiv(length, 64))].long() keys = page_bytes[page_ids, : 64 * 128].reshape(-1, 128) scales = page_bytes[page_ids, 64 * 128 :].contiguous().view(torch.float32).flatten() locs = locations[req, :length].long() @@ -335,11 +335,14 @@ def check(pages, block_table): graph = torch.cuda.CUDAGraph() with torch.cuda.graph(graph): pages, block_table = gather_paged_pools(packed, table, req_idx, lengths, max_pools) - # Empty HOLD rows must do no page writes, even with a 1M graph capacity. + # Empty rows initialize only their dummy page, even at 1M graph capacity. pages.fill_(127) lengths.zero_() graph.replay() - assert (pages == 127).all() + check(pages, block_table) + untouched = torch.ones(pages.shape[0], dtype=torch.bool, device=pages.device) + untouched[block_table[:, 0].long()] = False + assert (pages[untouched] == 127).all() for counts in ([513, 65, max_pools], [1, 0, 3]): lengths.copy_(torch.tensor(counts, device="cuda", dtype=torch.int32)) req_idx.copy_(torch.tensor([1, 3, 0], device="cuda", dtype=torch.int32)) From a55770c772991f6434310cc2280f22284e614b57 Mon Sep 17 00:00:00 2001 From: sufubao Date: Wed, 7 Oct 2026 22:41:23 +0800 Subject: [PATCH 08/40] Simplify ReplaySSM updates and cache metadata access --- .../triton_kernel/linear_att/replayssm.py | 55 +++++++------------ .../qwen3next_mem_manager.py | 8 +-- 2 files changed, 24 insertions(+), 39 deletions(-) diff --git a/lightllm/common/basemodel/triton_kernel/linear_att/replayssm.py b/lightllm/common/basemodel/triton_kernel/linear_att/replayssm.py index d68eedb22d..08b008e0a4 100644 --- a/lightllm/common/basemodel/triton_kernel/linear_att/replayssm.py +++ b/lightllm/common/basemodel/triton_kernel/linear_att/replayssm.py @@ -75,6 +75,21 @@ def _prepare_fold( tl.store(Active + N, tl.where(split, count, 0)) +@triton.jit +def _replay_update(state, raw_k, raw_v, g, beta, KDA: tl.constexpr): + raw_k /= tl.sqrt(tl.sum(raw_k * raw_k) + 1.0e-6) + decay = tl.exp(g) + if KDA: + state *= decay[:, None] + d = beta * (raw_v - tl.sum(state * raw_k[:, None], 0)) + state += raw_k[:, None] * d[None, :] + else: + sk = tl.sum(state * raw_k[:, None], 0) + d = beta * (raw_v - decay * sk) + state = state * decay + raw_k[:, None] * d[None, :] + return state + + @triton.jit def _fold_history( State, @@ -84,7 +99,6 @@ def _fold_history( Betas, Reqs, Cursors, - Cu, Active, N: tl.constexpr, LAYERS: tl.constexpr, @@ -93,10 +107,8 @@ def _fold_history( K: tl.constexpr, V: tl.constexpr, L: tl.constexpr, - HOLD: tl.constexpr, BK: tl.constexpr, BV: tl.constexpr, - VARLEN: tl.constexpr, SS: tl.constexpr, RS: tl.constexpr, VS: tl.constexpr, @@ -132,16 +144,7 @@ def _fold_history( else: g = tl.load(Gates + layer * GS + slot * (2 * L) + base + j).to(tl.float32) beta = tl.load(Betas + layer * BS + slot * (2 * L) + base + j).to(tl.float32) - raw_k /= tl.sqrt(tl.sum(raw_k * raw_k) + 1.0e-6) - decay = tl.exp(g) - if KDA: - state *= decay[:, None] - d = beta * (raw_v - tl.sum(state * raw_k[:, None], 0)) - state += raw_k[:, None] * d[None, :] - else: - sk = tl.sum(state * raw_k[:, None], 0) - d = beta * (raw_v - decay * sk) - state = state * decay + raw_k[:, None] * d[None, :] + state = _replay_update(state, raw_k, raw_v, g, beta, KDA) # Start the new history from the same rounded checkpoint future calls load. state = state.to(State.dtype.element_ty).to(tl.float32) tl.store(sp, state, (kk[:, None] < K) & (vv[None, :] < V)) @@ -179,7 +182,6 @@ def _replay( K: tl.constexpr, V: tl.constexpr, L: tl.constexpr, - WIDTH: tl.constexpr, HOLD: tl.constexpr, BK: tl.constexpr, BV: tl.constexpr, @@ -224,16 +226,7 @@ def _replay( else: 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) - decay = tl.exp(g) - if KDA: - state *= decay[:, None] - d = beta * (raw_v - tl.sum(state * raw_k[:, None], 0)) - state += raw_k[:, None] * d[None, :] - else: - sk = tl.sum(state * raw_k[:, None], 0) - d = beta * (raw_v - decay * sk) - state = state * decay + raw_k[:, None] * d[None, :] + state = _replay_update(state, raw_k, raw_v, g, beta, KDA) # Start the new history from the same rounded checkpoint future calls load. state = state.to(State.dtype.element_ty).to(tl.float32) tl.store(sp, state, (kk[:, None] < K) & (vv[None, :] < V)) @@ -617,7 +610,6 @@ def prepare_decode(self, reqs, cu_seqlens=None): self.betas, reqs, positions, - cu_seqlens, active, reqs.numel(), layers, @@ -626,10 +618,8 @@ def prepare_decode(self, reqs, cu_seqlens=None): kd, vd, self.capacity, - self.hold, triton.next_power_of_2(kd), 32, - cu_seqlens is not None, self.state.stride(0), self.raw_keys.stride(0), self.raw_values.stride(0), @@ -696,18 +686,16 @@ def snapshot_accepted_state(self, req_idx): self._materialize_accepted_state(reqs, output, snapshot=True) return output - def forward(self, layer, q, k, v, a, b, a_log, bias, reqs, cu_seqlens=None, run_config=None): - if layer == 0 and run_config is None: + def forward(self, layer, q, k, v, a, b, a_log, bias, reqs, cu_seqlens=None): + if layer == 0: configure_cache(self, "replay", q, k, v, a, b, a_log, bias, cu_seqlens) hv, kd, vd = self.state.shape[-3:] assert self.num_key_heads in (q.shape[-2], hv) axis = 1 if cu_seqlens is not None else 0 out = torch.empty_like(v) - config = run_config or self.run_config or {"BV": 32, "num_warps": 1} + config = self.run_config or {"BV": 32, "num_warps": 1} bv = config["BV"] - precompute_state = ( - self.verify_width > 1 and bv >= 16 and config.get("precompute_state", self.projection_mode == "precompute") - ) + precompute_state = self.verify_width > 1 and bv >= 16 and self.projection_mode == "precompute" _replay[(triton.cdiv(vd, bv), reqs.numel(), hv)]( q, k, @@ -738,7 +726,6 @@ def forward(self, layer, q, k, v, a, b, a_log, bias, reqs, cu_seqlens=None, run_ kd, vd, self.capacity, - self.verify_width, self.hold, triton.next_power_of_2(kd), bv, 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 452ed557ec..0e4e3364c9 100644 --- a/lightllm/common/kv_cache_mem_manager/qwen3next_mem_manager.py +++ b/lightllm/common/kv_cache_mem_manager/qwen3next_mem_manager.py @@ -267,9 +267,7 @@ def read_page_to_req( return def _get_req_state_indexes(self, req_idx: int): - 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 + return req_idx, req_idx * self.mem_manager.ssm_slots_per_req def _write_one_rank( self, @@ -279,7 +277,7 @@ def _write_one_rank( conv_page: torch.Tensor, ssm_page: torch.Tensor, ): - ssm_updates = getattr(mem, "ssm_update_cache", None) + ssm_updates = mem.ssm_update_cache 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 ssm_updates is not None and mem.req_to_mtp_state_index is not None: @@ -467,7 +465,7 @@ def _read_one_rank( conv_page: torch.Tensor, ssm_page: torch.Tensor, ): - ssm_updates = getattr(mem, "ssm_update_cache", None) + ssm_updates = mem.ssm_update_cache if ssm_updates is not None: ssm_updates.clear_history(req_idx) if mem.req_to_mtp_state_index is not None: From 3af33c2f47f1738e17eb70a3c1e8f42e25ed65ed Mon Sep 17 00:00:00 2001 From: sufubao Date: Thu, 8 Oct 2026 00:46:19 +0800 Subject: [PATCH 09/40] Tune SSM layouts per batch and preserve verify configs --- docs/replayssm/README.md | 10 +- .../common/basemodel/attention/linear/gdn.py | 27 ++-- .../common/basemodel/attention/linear/kda.py | 10 +- lightllm/common/basemodel/basemodel.py | 1 + lightllm/common/basemodel/batch_objs.py | 3 + lightllm/common/basemodel/infer_struct.py | 1 + .../triton_kernel/linear_att/replayssm.py | 10 +- .../linear_att/replayssm_compact.py | 25 +-- .../triton_kernel/linear_att/ssm_autotune.py | 45 +++--- lightllm/common/req_manager/hybrid_base.py | 4 +- lightllm/common/req_manager/linear_att.py | 16 +- .../mode_backend/chunked_prefill/impl.py | 1 + .../mode_backend/dp_backend/impl.py | 6 + .../model_infer/mtp_speculative/utils.py | 2 + .../linear_att/test_ssm_autotune.py | 144 +++++++++++++++++- unit_tests/models/test_post_layer_output.py | 2 +- 16 files changed, 240 insertions(+), 67 deletions(-) diff --git a/docs/replayssm/README.md b/docs/replayssm/README.md index f80c6333c1..dd88de64a9 100644 --- a/docs/replayssm/README.md +++ b/docs/replayssm/README.md @@ -51,10 +51,12 @@ promise bitwise equivalence or identical generated text to native BF16. The use an independent recurrence with each mode's rounding boundary. The SSM autotuner measures prepare, forward, acceptance and fold on disposable -state. Its selected layout is fixed before CUDA Graph capture and reused for -verification and checkpoint reconstruction; serving state is never a tuning -input. Cached configurations distinguish precision, stride, heads, verify width, -history capacity, projection mode and recurrent rule. +state; serving state is never a tuning input. Each batch bucket selects its own +layout before CUDA Graph capture. Model outputs carry that layout through +unpadding, and compact acceptance uses the corresponding verify layout for each +microbatch. Replay checkpoint reconstruction keeps a batch-independent layout. +Cached configurations distinguish token count, sequence capacity, precision, +stride, heads, verify width, history capacity, projection mode and recurrent rule. KDA uses a per-key-dimension log decay, while GDN uses one scalar per value head. KDA history stores the full gate vector; checkpoint projection applies diff --git a/lightllm/common/basemodel/attention/linear/gdn.py b/lightllm/common/basemodel/attention/linear/gdn.py index 89981ad934..5f1c988e45 100644 --- a/lightllm/common/basemodel/attention/linear/gdn.py +++ b/lightllm/common/basemodel/attention/linear/gdn.py @@ -357,18 +357,11 @@ def _gdn_decode_kernel( layer_weight.layer_num_ - layer_weight.layer_num_ // backend.req_manager.linear_config.full_attention_interval ) + args = (query, key, value, a, b, layer_weight.linear_A_log.weight, layer_weight.linear_dt_bias.weight) + if layer == 0: + self.infer_state.ssm_run_config = ssm_updates.get_run_config(*args) return ( - ssm_updates.forward( - layer, - query, - key, - value, - a, - b, - layer_weight.linear_A_log.weight, - layer_weight.linear_dt_bias.weight, - self.b_conv_buffer_idx, - ), + ssm_updates.forward(layer, *args, self.b_conv_buffer_idx, run_config=self.infer_state.ssm_run_config), z, ) core_attn_out, _ = fused_recurrent_gated_delta_rule( @@ -422,17 +415,15 @@ def _gdn_mtp_kernel( layer_weight.layer_num_ - layer_weight.layer_num_ // backend.req_manager.linear_config.full_attention_interval ) + args = (query, key, value, a, b, layer_weight.linear_A_log.weight, layer_weight.linear_dt_bias.weight) + if layer == 0: + self.infer_state.ssm_run_config = ssm_updates.get_run_config(*args, cu_seqlens_q) return ssm_updates.forward( layer, - query, - key, - value, - a, - b, - layer_weight.linear_A_log.weight, - layer_weight.linear_dt_bias.weight, + *args, self.b_conv_buffer_idx, cu_seqlens_q, + run_config=self.infer_state.ssm_run_config, ) 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, diff --git a/lightllm/common/basemodel/attention/linear/kda.py b/lightllm/common/basemodel/attention/linear/kda.py index a90c0e42ef..6b36311879 100644 --- a/lightllm/common/basemodel/attention/linear/kda.py +++ b/lightllm/common/basemodel/attention/linear/kda.py @@ -237,8 +237,7 @@ def decode_att( updates = self.infer_state.req_manager.ssm_update_cache if updates is not None: layer = layer_num - layer_num // self.infer_state.req_manager.linear_config.full_attention_interval - return updates.forward( - layer, + args = ( q, k, v, @@ -246,8 +245,15 @@ def decode_att( raw_beta.view(-1, backend.tp_num_heads), layer_weight.linear_A_log.weight, layer_weight.linear_dt_bias.weight, + ) + if layer == 0: + self.infer_state.ssm_run_config = updates.get_run_config(*args, self.b1_mtp_cu_q_seq_len) + return updates.forward( + layer, + *args, self.b_conv_buffer_idx, self.b1_mtp_cu_q_seq_len, + run_config=self.infer_state.ssm_run_config, ) output, _ = fused_recurrent_kda( q=q, diff --git a/lightllm/common/basemodel/basemodel.py b/lightllm/common/basemodel/basemodel.py index 5a2d5b77d8..02478d3e0b 100755 --- a/lightllm/common/basemodel/basemodel.py +++ b/lightllm/common/basemodel/basemodel.py @@ -491,6 +491,7 @@ def _create_model_output(self, post_output: PostLayerOutput, infer_state: InferS mtp_collector=infer_state.hidden_collector.finish_output(infer_state=infer_state), prompt_logics=infer_state.prompt_logics, logits_token_ids=post_output.logits_token_ids, + ssm_run_config=infer_state.ssm_run_config, ) return output diff --git a/lightllm/common/basemodel/batch_objs.py b/lightllm/common/basemodel/batch_objs.py index d8f2329b24..a393c5c3e0 100644 --- a/lightllm/common/basemodel/batch_objs.py +++ b/lightllm/common/basemodel/batch_objs.py @@ -202,6 +202,9 @@ class ModelOutput: # target 模型会将候选值回填到完整词表 logits，因此不需要该映射。 logits_token_ids: Optional[torch.Tensor] = None + # Layout selected by this verify graph; preserved when logits are unpadded. + ssm_run_config: Optional[dict] = None + def __post_init__(self) -> None: if self.mtp_collector is None: self.mtp_collector = ModelMtpOutputCollector() diff --git a/lightllm/common/basemodel/infer_struct.py b/lightllm/common/basemodel/infer_struct.py index 0605478e35..26ef020348 100755 --- a/lightllm/common/basemodel/infer_struct.py +++ b/lightllm/common/basemodel/infer_struct.py @@ -64,6 +64,7 @@ def __init__(self): # 在microbatch overlap的运行模式下，用于标记当前 microbatch 的 index 序号 # 在一些细节场景下需要有该信息区分一些资源的申请和管理。 self.microbatch_index: int = 0 + self.ssm_run_config: Optional[dict] = None # 当前 forward 独占的 hidden state 收集器。普通推理使用短生命周期实例； # Prefill CUDA Graph 使用随 graph infer state 长期保存的实例。 diff --git a/lightllm/common/basemodel/triton_kernel/linear_att/replayssm.py b/lightllm/common/basemodel/triton_kernel/linear_att/replayssm.py index 08b008e0a4..82eb65e67f 100644 --- a/lightllm/common/basemodel/triton_kernel/linear_att/replayssm.py +++ b/lightllm/common/basemodel/triton_kernel/linear_att/replayssm.py @@ -547,6 +547,7 @@ def __init__( self.projection_mode = projection_mode self.run_config = run_config self._config_is_fixed = run_config is not None + self.batch_configs = {} layers, slots, hv, k, v = state.shape # Omission preserves the redundant layout for existing callers and A/B. self.num_key_heads = hv if num_key_heads is None else num_key_heads @@ -686,14 +687,15 @@ def snapshot_accepted_state(self, req_idx): self._materialize_accepted_state(reqs, output, snapshot=True) return output - def forward(self, layer, q, k, v, a, b, a_log, bias, reqs, cu_seqlens=None): - if layer == 0: - configure_cache(self, "replay", q, k, v, a, b, a_log, bias, cu_seqlens) + def get_run_config(self, q, k, v, a, b, a_log, bias, cu_seqlens=None): + return configure_cache(self, "replay", q, k, v, a, b, a_log, bias, cu_seqlens) + + def forward(self, layer, q, k, v, a, b, a_log, bias, reqs, cu_seqlens=None, run_config=None): + config = run_config or self.get_run_config(q, k, v, a, b, a_log, bias, cu_seqlens) hv, kd, vd = self.state.shape[-3:] assert self.num_key_heads in (q.shape[-2], hv) axis = 1 if cu_seqlens is not None else 0 out = torch.empty_like(v) - config = self.run_config or {"BV": 32, "num_warps": 1} bv = config["BV"] precompute_state = self.verify_width > 1 and bv >= 16 and self.projection_mode == "precompute" _replay[(triton.cdiv(vd, bv), reqs.numel(), hv)]( diff --git a/lightllm/common/basemodel/triton_kernel/linear_att/replayssm_compact.py b/lightllm/common/basemodel/triton_kernel/linear_att/replayssm_compact.py index 6799be5f34..8f55269fd9 100644 --- a/lightllm/common/basemodel/triton_kernel/linear_att/replayssm_compact.py +++ b/lightllm/common/basemodel/triton_kernel/linear_att/replayssm_compact.py @@ -125,6 +125,7 @@ def __init__( self.verify_width = verify_width # Verify and accepted-state reconstruction must share the same layout. self._config_is_fixed = run_config is not None + self.batch_configs = {} self.run_config = dict(run_config or {"BV": 8, "num_warps": 1}) assert self.run_config["BV"] in (8, 16, 32, 64, 128) assert self.run_config["num_warps"] in (1, 2, 4, 8) @@ -155,15 +156,17 @@ def merge_accepted_updates(self, reqs): def snapshot_accepted_state(self, req_idx): return self.state[:, req_idx] - def forward(self, layer, q, k, v, a, b, a_log, bias, reqs, cu_seqlens=None): + def get_run_config(self, q, k, v, a, b, a_log, bias, cu_seqlens=None): + return configure_cache(self, "gdn", q, k, v, a, b, a_log, bias, cu_seqlens) + + def forward(self, layer, q, k, v, a, b, a_log, bias, reqs, cu_seqlens=None, run_config=None): assert cu_seqlens is not None, "compact replay is only used for speculative verify" - if layer == 0: - configure_cache(self, "gdn", q, k, v, a, b, a_log, bias, cu_seqlens) + config = run_config or self.get_run_config(q, k, v, a, b, a_log, bias, cu_seqlens) _, slots, hv, kd, vd = self.state.shape assert self.num_key_heads in (q.shape[-2], hv) out = torch.empty_like(v) # Match native reduction layouts: changing BV can alter BF16 rounding. - bv = self.run_config["BV"] + bv = config["BV"] _compact[(triton.cdiv(vd, bv), reqs.numel(), hv)]( q, k, @@ -199,15 +202,17 @@ def forward(self, layer, q, k, v, a, b, a_log, bias, reqs, cu_seqlens=None): False, self.kda, self.lower_bound, - num_warps=self.run_config["num_warps"], - num_stages=self.run_config.get("num_stages", 3), + num_warps=config["num_warps"], + num_stages=config.get("num_stages", 3), ) return out - def accept_updates(self, reqs, accepted): + def accept_updates(self, reqs, accepted, run_config=None): """Merge this round's prefix; accepted contains per-request last accepted indexes.""" + assert run_config is not None or not self.batch_configs, "compact acceptance requires the verify configuration" layers, slots, hv, kd, vd = self.state.shape - bv = self.run_config["BV"] + config = run_config or self.run_config + bv = config["BV"] _compact[(triton.cdiv(vd, bv), reqs.numel(), layers * hv)]( None, None, @@ -243,6 +248,6 @@ def accept_updates(self, reqs, accepted): True, self.kda, self.lower_bound, - num_warps=self.run_config["num_warps"], - num_stages=self.run_config.get("num_stages", 3), + num_warps=config["num_warps"], + num_stages=config.get("num_stages", 3), ) diff --git a/lightllm/common/basemodel/triton_kernel/linear_att/ssm_autotune.py b/lightllm/common/basemodel/triton_kernel/linear_att/ssm_autotune.py index 032c7daeb7..f3d5ab06ac 100644 --- a/lightllm/common/basemodel/triton_kernel/linear_att/ssm_autotune.py +++ b/lightllm/common/basemodel/triton_kernel/linear_att/ssm_autotune.py @@ -1,5 +1,6 @@ """Tune a complete SSM update cycle without touching the serving state pool.""" +import ast import torch import triton @@ -24,6 +25,7 @@ def static_key(cache, mode, q, k, v, a, b, cu_seqlens): # state precision, recurrent rule and history capacity. key = { "mode": mode, + "v": 5, "shape": str((layers, q.shape[-2], hv, cache.num_key_heads, kd, vd, cache.verify_width)), "dtype": str((q.dtype, cache.state.dtype)), "strides": str((q.stride(axis), k.stride(axis), v.stride(axis), a.stride(0), b.stride(0))), @@ -35,7 +37,6 @@ def static_key(cache, mode, q, k, v, a, b, cu_seqlens): key["lower"] = cache.lower_bound if mode == "replay": key["history_dtype"] = str(cache.keys.dtype) - key["v"] = 4 # Small batches omit the forward fold already handled by prepare. key["projection"] = ( ("kda_mma" if cache.kda else "block_mma_activation_precision") if cache.projection_mode == "precompute" and cache.verify_width > 1 @@ -46,11 +47,12 @@ def static_key(cache, mode, q, k, v, a, b, cu_seqlens): def run_key(cache, q, cu_seqlens): tokens = q.shape[1 if cu_seqlens is not None else 0] - return triton.cdiv(tokens, cache.verify_width) + return tokens, cu_seqlens.numel() - 1 if cu_seqlens is not None else tokens def rebuild_inputs(cache, mode, q, k, v, a, b, a_log, bias, cu_seqlens=None, workload=None, run_config=None): - batch = run_key(cache, q, cu_seqlens) + tokens, sequences = run_key(cache, q, cu_seqlens) + batch = min(sequences, triton.cdiv(tokens, cache.verify_width)) if cu_seqlens is not None else tokens layers, _, hv, kd, vd = cache.state.shape state = torch.empty((layers, batch + 1, hv, kd, vd), device=q.device, dtype=cache.state.dtype).normal_(0, 0.02) if mode == "replay": @@ -79,13 +81,17 @@ def rebuild_inputs(cache, mode, q, k, v, a, b, a_log, bias, cu_seqlens=None, wor torch.empty_strided(x.shape, x.stride(), device=x.device, dtype=x.dtype).normal_(0, 0.2) for x in (q, k, v, a, b, a_log, bias) ] - reqs = torch.arange(batch, dtype=torch.int32, device=q.device) - tokens = q.shape[1 if cu_seqlens is not None else 0] - cu = (torch.arange(batch + 1, dtype=torch.int32, device=q.device) * cache.verify_width).clamp_max(tokens) - lengths = cu[1:] - cu[:-1] + reqs = torch.cat( + ( + torch.arange(batch, dtype=torch.int32, device=q.device), + torch.full((sequences - batch,), batch, dtype=torch.int32, device=q.device), + ) + ) + cu = (torch.arange(sequences + 1, dtype=torch.int32, device=q.device) * cache.verify_width).clamp_max(tokens) + lengths = cu[1 : batch + 1] - cu[:batch] accepted = torch.cat((lengths - 1, lengths.new_zeros(1))) partial = accepted.clone() - partial[:-1] = torch.minimum(reqs % cache.verify_width, lengths - 1) + partial[:-1] = torch.minimum(reqs[:batch] % cache.verify_width, lengths - 1) snapshot = torch.empty((layers, hv, kd, vd), device=q.device, dtype=state.dtype) return (), dict( cache=scratch, @@ -107,6 +113,7 @@ def rebuild_inputs(cache, mode, q, k, v, a, b, a_log, bias, cu_seqlens=None, wor configs_gen_func=get_configs, static_key_func=static_key, run_key_func=run_key, + run_key_distance_func=lambda a, b: sum(abs(x - y) for x, y in zip(ast.literal_eval(a), ast.literal_eval(b))), kernel_type=AutotuneKernelType.DECODE_ATTENTION, rebuild_input_func=rebuild_inputs, warmup_all_exist_config=False, @@ -138,13 +145,17 @@ def select_config(cache, mode, q, k, v, a, b, a_log, bias, cu_seqlens=None, work def configure_cache(cache, mode, q, k, v, a, b, a_log, bias, cu_seqlens): - if cache._config_is_fixed or not Autotuner.is_kernel_autotune_warmup(AutotuneKernelType.DECODE_ATTENTION): - return + default = cache.run_config or {"BV": 32, "num_warps": 1} + if cache._config_is_fixed: + return default + key = (tuple(static_key(cache, mode, q, k, v, a, b, cu_seqlens).items()), run_key(cache, q, cu_seqlens)) + if key in cache.batch_configs: + return cache.batch_configs[key] + if not Autotuner.is_kernel_autotune_warmup(AutotuneKernelType.DECODE_ATTENTION): + return default config = select_config(cache, mode, q, k, v, a, b, a_log, bias, cu_seqlens) - if config is not None: - cache.run_config = config - logger.info(f"SSM {mode} update-cycle config: {config}, warmup requests: {run_key(cache, q, cu_seqlens)}") - # One layout per cache, selected at the largest graph warmup. - # Per-batch tuning needs explicit layout propagation through verify/accept, - # including graph padding and overlapping microbatches; never use "last config". - cache._config_is_fixed = True + config = dict(config or default) + # Each captured graph keeps its own layout, even after other buckets warm up. + cache.batch_configs[key] = config + logger.info(f"SSM {mode} update-cycle config: {config}, warmup batch: {key[1]}") + return config diff --git a/lightllm/common/req_manager/hybrid_base.py b/lightllm/common/req_manager/hybrid_base.py index 948313116d..e71e86edb5 100644 --- a/lightllm/common/req_manager/hybrid_base.py +++ b/lightllm/common/req_manager/hybrid_base.py @@ -65,6 +65,8 @@ def save_big_page_states(self, b_req_idx: torch.Tensor, req_indexes: List[int], def save_state(self, req_idx: int, buffer_idx: int, state_cache_manager): """请求 GPU 运行态 → 指定 CPU checkpoint 槽位；大小页共用，调用方负责分配槽位。""" - def update_mtp_state(self, b_req_mtp_start_loc, b_req_idx, b_mtp_index, accepted_index, verify_width): + def update_mtp_state( + self, b_req_mtp_start_loc, b_req_idx, b_mtp_index, accepted_index, verify_width, ssm_accept_batches=None + ): """接受推测 token 后更新运行态位置；需要调整状态索引的模型覆写。""" return diff --git a/lightllm/common/req_manager/linear_att.py b/lightllm/common/req_manager/linear_att.py index 8c796daa84..1914421e42 100644 --- a/lightllm/common/req_manager/linear_att.py +++ b/lightllm/common/req_manager/linear_att.py @@ -143,7 +143,9 @@ def get_mamba_cache(self, layer_idx_in_all: int): ssm_states = self.req_to_ssm_state.buffer[layer_idx_in_linear] return conv_states, ssm_states - def update_mtp_state(self, b_req_mtp_start_loc, b_req_idx, b_mtp_index, accepted_index, verify_width): + def update_mtp_state( + self, b_req_mtp_start_loc, b_req_idx, b_mtp_index, accepted_index, verify_width, ssm_accept_batches=None + ): from lightllm.common.basemodel.triton_kernel.mtp_utils import linear_att_mtp_state_index_update replay_cursors = None @@ -165,7 +167,17 @@ def update_mtp_state(self, b_req_mtp_start_loc, b_req_idx, b_mtp_index, accepted if self.ssm_update_cache is not None and replay_cursors is None: reqs = b_req_idx[b_req_mtp_start_loc.long()] - self.ssm_update_cache.accept_updates(reqs, self.req_to_mtp_state_index) + if ssm_accept_batches is None or isinstance(self.ssm_update_cache, ReplaySSMCache): + self.ssm_update_cache.accept_updates(reqs, self.req_to_mtp_state_index) + else: + assert sum(size for size, _ in ssm_accept_batches) == reqs.numel() + start = 0 + for size, config in ssm_accept_batches: + if size: + self.ssm_update_cache.accept_updates( + reqs[start : start + size], self.req_to_mtp_state_index, run_config=config + ) + start += size def restore_state(self, req: "InferReq", state_cache_manager: LinearAttCacheManager, buffer_idx: int): if self.ssm_update_cache is not None: diff --git a/lightllm/server/router/model_infer/mode_backend/chunked_prefill/impl.py b/lightllm/server/router/model_infer/mode_backend/chunked_prefill/impl.py index d0026da9bd..e33a22ace0 100644 --- a/lightllm/server/router/model_infer/mode_backend/chunked_prefill/impl.py +++ b/lightllm/server/router/model_infer/mode_backend/chunked_prefill/impl.py @@ -284,6 +284,7 @@ def decode_mtp( b_req_idx=model_input.b_req_idx, b_req_mtp_start_loc=b_req_mtp_start_loc, b_mtp_index=model_input.b_mtp_index, + ssm_accept_batches=((req_num, model_output.ssm_run_config),), ) accepted_index_cpu = g_pin_mem_manager.async_copy_from_gpu_tensor( key="accepted_index", diff --git a/lightllm/server/router/model_infer/mode_backend/dp_backend/impl.py b/lightllm/server/router/model_infer/mode_backend/dp_backend/impl.py index 2cb4e89e5b..2ef4770424 100644 --- a/lightllm/server/router/model_infer/mode_backend/dp_backend/impl.py +++ b/lightllm/server/router/model_infer/mode_backend/dp_backend/impl.py @@ -515,6 +515,7 @@ def decode_mtp(self, event_pack: OverlapEventPack, decode_reqs: List[InferReq]): b_req_idx=model_input.b_req_idx, b_req_mtp_start_loc=b_req_mtp_start_loc, b_mtp_index=model_input.b_mtp_index, + ssm_accept_batches=((req_num, model_output.ssm_run_config),), ) accepted_index_cpu = g_pin_mem_manager.async_copy_from_gpu_tensor( key="accepted_index", @@ -576,6 +577,7 @@ def decode_mtp(self, event_pack: OverlapEventPack, decode_reqs: List[InferReq]): sync_event.record() if req_num > 0: + # 第二阶段 event_pack.notify_post_handle_and_wait_pre_post_handle() if spec_plan.skip_verify_sync: @@ -788,6 +790,10 @@ def decode_overlap_mtp(self, event_pack: OverlapEventPack, decode_reqs: List[Inf b_req_idx=b_req_idx, b_req_mtp_start_loc=b_req_mtp_start_loc, b_mtp_index=b_mtp_index, + ssm_accept_batches=( + (real_request_num0, model_output0.ssm_run_config), + (real_request_num1, model_output1.ssm_run_config), + ), ) mtp_accept_len0 = mtp_accept_len[:real_request_num0] mtp_accept_len1 = mtp_accept_len[real_request_num0:] diff --git a/lightllm/server/router/model_infer/mtp_speculative/utils.py b/lightllm/server/router/model_infer/mtp_speculative/utils.py index 7835313bb0..366668530f 100644 --- a/lightllm/server/router/model_infer/mtp_speculative/utils.py +++ b/lightllm/server/router/model_infer/mtp_speculative/utils.py @@ -22,6 +22,7 @@ def verify_mtp_tokens( b_req_idx: torch.Tensor, b_req_mtp_start_loc: torch.Tensor, b_mtp_index: torch.Tensor, + ssm_accept_batches=None, ) -> Tuple[torch.Tensor, torch.Tensor]: """Verify target tokens and update recurrent MTP state when required.""" @@ -38,6 +39,7 @@ def verify_mtp_tokens( b_mtp_index=b_mtp_index, accepted_index=accepted_index, verify_width=backend.max_draft_step + 1, + ssm_accept_batches=ssm_accept_batches, ) return accept_lengths, accepted_index diff --git a/unit_tests/common/basemodel/triton_kernel/linear_att/test_ssm_autotune.py b/unit_tests/common/basemodel/triton_kernel/linear_att/test_ssm_autotune.py index 396c3e863d..0665b949ac 100644 --- a/unit_tests/common/basemodel/triton_kernel/linear_att/test_ssm_autotune.py +++ b/unit_tests/common/basemodel/triton_kernel/linear_att/test_ssm_autotune.py @@ -125,10 +125,10 @@ def test_layout_is_fixed_before_graph_capture_and_cached_configs_do_not_execute( tuning.configure_cache(**inputs) assert cache.run_config == default and not cache._config_is_fixed with Autotuner.autotune_warmup(AutotuneKernelType.DECODE_ATTENTION): - tuning.configure_cache(**inputs) - assert cache.run_config == (default if level == 3 else config) - assert cache._config_is_fixed - # Later/smaller graphs, captures, eager calls and acceptance share this layout. + selected = tuning.configure_cache(**inputs) + assert selected == (default if level == 3 else config) + assert cache.run_config == default and not cache._config_is_fixed + # Reusing the same bucket must not tune again. monkeypatch.setattr(tuning, "select_config", lambda *args, **kwargs: pytest.fail("layout changed after capture")) with Autotuner.autotune_warmup(AutotuneKernelType.DECODE_ATTENTION): tuning.configure_cache(**inputs) @@ -181,19 +181,26 @@ def checked_bench(*args, **kwargs): # Exercise the real forward hook with the HOLD-only graph warmup input. cache.forward(0, *args, reqs, inputs["cu_seqlens"]) assert len(timings) == len(configs) - assert cache.run_config in configs + selected = cache.get_run_config(*args, inputs["cu_seqlens"]) + assert selected in configs # An explicit identical layout supplies a reference for capture/replay and # alternating real/HOLD rows. It must never trigger its own tuning. expected = make_inputs(mode, dtype, "cuda", width=width, projection_mode=projection_mode, kda=kda)["cache"] - expected.run_config = cache.run_config.copy() + expected.run_config = cache.run_config expected._config_is_fixed = True expected.state.copy_(before["state"]) accepted = torch.zeros(cache.state.shape[1], device="cuda", dtype=torch.int32) def step(target): target.prepare_decode(reqs, inputs["cu_seqlens"]) - out = [target.forward(layer, *args, reqs, inputs["cu_seqlens"]) for layer in range(target.state.shape[0])] - target.accept_updates(reqs, accepted) + out = [ + target.forward(layer, *args, reqs, inputs["cu_seqlens"], run_config=selected) + for layer in range(target.state.shape[0]) + ] + if mode == "gdn": + target.accept_updates(reqs, accepted, run_config=selected) + else: + target.accept_updates(reqs, accepted) return out reqs.fill_(cache.hold) @@ -261,3 +268,124 @@ def test_kda_production_tuning_cache_can_be_written(monkeypatch, tmp_path, tp, d path = tmp_path / KernelConfigs.get_config_file_name(key) path.write_text("{}") assert path.is_file() + + +def test_batch_configs_survive_other_warmups_and_preserve_padding(monkeypatch): + inputs = make_inputs("gdn") + cache = inputs["cache"] + small = dict(inputs, **{name: inputs[name][:, :4] for name in ("q", "k", "v")}) + small["a"], small["b"] = inputs["a"][:4], inputs["b"][:4] + small["cu_seqlens"] = torch.tensor([0, 4, 4], dtype=torch.int32) + selected = [] + + def choose(cache, mode, q, *args): + config = {"BV": 16 if q.shape[1] == 4 else 32, "num_warps": 2} + selected.append(config) + return config + + monkeypatch.setattr(tuning, "select_config", choose) + with Autotuner.autotune_warmup(AutotuneKernelType.DECODE_ATTENTION): + large_config = tuning.configure_cache(**inputs) + small_config = tuning.configure_cache(**small) + assert tuning.configure_cache(**inputs) is large_config + assert len(selected) == 2 and large_config != small_config + assert tuning.configure_cache(**small) is small_config + assert cache.run_config == {"BV": 8, "num_warps": 1} + _, rebuilt = tuning.rebuild_inputs(**small) + assert rebuilt["workload"][0].tolist() == [0, rebuilt["cache"].hold] + assert rebuilt["cu_seqlens"].tolist() == [0, 4, 4] + assert tuning.run_key(**{k: rebuilt[k] for k in ("cache", "q", "cu_seqlens")}) == (4, 2) + # A graph token bucket can also exceed the request-pool capacity. + _, limited = tuning.rebuild_inputs(**dict(inputs, cu_seqlens=torch.tensor([0, 4, 8], dtype=torch.int32))) + assert limited["workload"][0].tolist() == [0, 1] + assert limited["cu_seqlens"].tolist() == [0, 4, 8] + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") +@pytest.mark.parametrize("mode,projection", [("gdn", "inline"), ("replay", "inline"), ("replay", "precompute")]) +@pytest.mark.parametrize("dtype", [torch.float32, torch.bfloat16]) +@pytest.mark.parametrize("kda", [False, True]) +@pytest.mark.parametrize("dynamic", [False, True]) +def test_two_verify_graphs_accept_with_their_own_configs(monkeypatch, mode, projection, dtype, kda, dynamic): + from types import SimpleNamespace + from lightllm.common.basemodel.basemodel import TpPartBaseModel + from lightllm.common.basemodel.batch_objs import PostLayerOutput, ModelMtpOutputCollector + from lightllm.common.req_manager.linear_att import ReqManagerForMamba + + inputs = make_inputs(mode, dtype, "cuda", projection_mode=projection, kda=kda) + cache = inputs["cache"] + initial = cache.state.clone() + configs = [{"BV": 8, "num_warps": 1, "num_stages": 1}, {"BV": 32, "num_warps": 2, "num_stages": 1}] + monkeypatch.setattr(tuning, "select_config", lambda *args: configs[args[2].shape[1] > 8]) + manager = object.__new__(ReqManagerForMamba) + manager.ssm_update_cache = cache + manager.req_to_mtp_state_index = torch.zeros(9, dtype=torch.int32, device="cuda") + buckets = [] + for index, ids in enumerate(([0], [1, 3])): + # The physical verify buckets include HOLD, but acceptance is unpadded. + count = len(ids) + 1 + args = [inputs[name] for name in ("q", "k", "v", "a", "b", "a_log", "bias")] + args[:3] = [x[:, : count * 4].contiguous() for x in args[:3]] + args[3:5] = [x[: count * 4].contiguous() for x in args[3:5]] + capacity = min(count * 4, cache.hold) if dynamic else count + reqs = torch.tensor(ids + [cache.hold] * (capacity - len(ids)), dtype=torch.int32, device="cuda") + cu = (torch.arange(capacity + 1, dtype=torch.int32, device="cuda") * 4).clamp_max(count * 4) + with Autotuner.autotune_warmup(AutotuneKernelType.DECODE_ATTENTION): + config = cache.get_run_config(*args, cu) + assert config == configs[index] + ref = make_inputs(mode, dtype, "cuda", projection_mode=projection, kda=kda)["cache"] + ref.state.copy_(initial) + ref.run_config = config + ref._config_is_fixed = True + + def forward(target): + target.prepare_decode(reqs, cu) + return [target.forward(layer, *args, reqs, cu, run_config=config) for layer in range(2)] + + forward(cache) + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + outputs = forward(cache) + state = SimpleNamespace( + hidden_collector=SimpleNamespace(finish_output=lambda **kw: ModelMtpOutputCollector()), + prompt_logics=None, + ssm_run_config=config, + ) + output = TpPartBaseModel._create_model_output( + None, PostLayerOutput(logits=outputs[0].view(count * 4, -1)), state + ) + output = TpPartBaseModel._create_unpad_decode_model_output(None, output, len(ids) * 4) + assert output.ssm_run_config is config + buckets.append((graph, outputs, output, ref, args, reqs, cu, ids)) + cache.state.copy_(initial) + if mode == "replay": + cache.cursors.zero_() + req_order = [0, 1, 3] + rows = torch.tensor([r for r in req_order for _ in range(4)], dtype=torch.int32, device="cuda") + mtp = torch.arange(4, dtype=torch.int32, device="cuda").repeat(3) + starts = torch.arange(3, dtype=torch.int32, device="cuda") * 4 + for iteration in range(24): + # Alternate replay order; neither Graph replay executes Python config selection. + for bucket_index in [0, 1] if iteration % 2 else [1, 0]: + graph, outputs, output, ref, args, reqs, cu, ids = buckets[bucket_index] + graph.replay() + ref.prepare_decode(reqs, cu) + for layer in range(2): + expected = ref.forward(layer, *args, reqs, cu) + torch.testing.assert_close(outputs[layer], expected, atol=0, rtol=0) + counts = [1 + (iteration + i) % 4 for i in range(3)] + flags = torch.tensor([int(j < n) for n in counts for j in range(4)], dtype=torch.int32, device="cuda") + manager.update_mtp_state( + starts, + rows, + mtp, + flags, + 4, + ssm_accept_batches=tuple((len(b[-1]), b[2].ssm_run_config) for b in buckets), + ) + for _, _, _, ref, _, reqs, _, ids in buckets: + ref.accept_updates(reqs[:-1], manager.req_to_mtp_state_index) + torch.testing.assert_close(cache.state[:, ids], ref.state[:, ids], atol=0, rtol=0) + if mode == "replay": + assert torch.equal(cache.cursors[ids], ref.cursors[ids]) + assert len(cache.batch_configs) == 2 diff --git a/unit_tests/models/test_post_layer_output.py b/unit_tests/models/test_post_layer_output.py index e020c6709b..e21d0b1d1c 100644 --- a/unit_tests/models/test_post_layer_output.py +++ b/unit_tests/models/test_post_layer_output.py @@ -290,7 +290,7 @@ def test_model_output_and_unpadding_preserve_candidate_mapping() -> None: model = TpPartBaseModel.__new__(TpPartBaseModel) logits = torch.randn(3, 2) ids = torch.tensor([[10, 20], [30, 40], [50, 60]]) - state = SimpleNamespace(hidden_collector=NoopHiddenCollector(), prompt_logics=None) + state = SimpleNamespace(hidden_collector=NoopHiddenCollector(), prompt_logics=None, ssm_run_config=None) output = model._create_model_output(PostLayerOutput(logits, ids), state) unpadded = model._create_unpad_decode_model_output(output, 2) torch.testing.assert_close(unpadded.logits, logits[:2]) From 3fa7c6a28fb77a11f64bbc40f754401d1054f9f3 Mon Sep 17 00:00:00 2001 From: sufubao Date: Wed, 14 Oct 2026 11:08:58 +0800 Subject: [PATCH 10/40] perf(replayssm): optimize GDN MTP decoding and sampling --- docs/replayssm/README.md | 32 +++- docs/replayssm/h200_optimizations.md | 125 +++++++++++++++ docs/replayssm/h200_retest_summary.json | 86 +++++++++++ lightllm/common/basemodel/basemodel.py | 49 +++++- lightllm/common/basemodel/batch_objs.py | 7 + lightllm/common/basemodel/infer_struct.py | 2 + .../triton_kernel/linear_att/replayssm.py | 125 ++++++++++++--- .../triton_kernel/linear_att/ssm_autotune.py | 5 + .../post_process/apply_penalty_gpu_cache.py | 40 ++--- .../triton_kernel/post_process/argmax.py | 47 ++++++ .../triton_kernel/post_process/greedy.py | 66 ++++++++ lightllm/common/req_manager/linear_att.py | 1 + .../llama/layer_infer/post_layer_infer.py | 16 +- .../layer_infer/transformer_layer_infer.py | 2 +- lightllm/server/api_cli.py | 11 +- lightllm/server/api_start.py | 11 +- lightllm/server/core/objs/start_args_type.py | 1 + .../model_infer/mode_backend/base_backend.py | 11 ++ .../mode_backend/generic_post_process.py | 13 ++ .../proposers/eagle_with_att.py | 17 ++- .../kernels/benchmark_replayssm_service.py | 118 +++++++++++++++ .../common/basemodel/test_model_output.py | 35 +++++ .../linear_att/test_replayssm_parallel.py | 71 +++++++++ .../linear_att/test_ssm_autotune.py | 1 + .../triton_kernel/test_greedy_sampling.py | 142 ++++++++++++++++++ .../mtp_speculative/test_eagle_with_att.py | 47 +++++- unit_tests/server/test_ssm_state_mode.py | 27 ++++ 27 files changed, 1050 insertions(+), 58 deletions(-) create mode 100644 docs/replayssm/h200_optimizations.md create mode 100644 docs/replayssm/h200_retest_summary.json create mode 100644 lightllm/common/basemodel/triton_kernel/post_process/argmax.py create mode 100644 lightllm/common/basemodel/triton_kernel/post_process/greedy.py create mode 100644 test/benchmark/kernels/benchmark_replayssm_service.py create mode 100644 unit_tests/common/basemodel/triton_kernel/linear_att/test_replayssm_parallel.py create mode 100644 unit_tests/common/basemodel/triton_kernel/test_greedy_sampling.py diff --git a/docs/replayssm/README.md b/docs/replayssm/README.md index dd88de64a9..27cb5e5678 100644 --- a/docs/replayssm/README.md +++ b/docs/replayssm/README.md @@ -27,10 +27,37 @@ capacities are 4, 8, 16, 32 and 64. Non-native modes support Qwen3-Next/Qwen3.5 GDN and GLM-5.3-FLASH KDA with FP32 or BF16 state. Other models retain their native path. `--replayssm_projection_mode` -selects `inline` (default) or `precompute`; the latter is only valid with replay. +selects `inline` (default), `precompute`, or `parallel`; the latter two require replay. +`parallel` uses a Tensor Core triangular solve for BF16 GDN MTP with 128x128 +heads. It currently uses the measured BV128 / 2-warp configuration. Choose the projection mode and capacity using the target workload. Without MTP, replay adds history storage to the single native checkpoint. +### Qwen3.5 H200 experimental optimization + +For the measured Qwen3.5-27B / TP1 / BF16 / MTP4 configuration: + +```bash +--ssm_state_mode replay --replayssm_cache_len 16 \ +--replayssm_projection_mode parallel --replayssm_fold_method derived +``` + +`--replayssm_fold_method raw` remains the default. `derived` requires `parallel` +and folds cached FP32 keys/deltas with FP32 matrix accumulation. It changes +checkpoint rounding relative to raw-input replay and is not bitwise equivalent. +Neither mode changes the configured request, KV, or history capacity. + +For Qwen3.5 dense TP1 with fixed EAGLE-with-attention MTP, `parallel` also +captures first-draft graphs above the logical request capacity and projects only +accepted-tail hidden rows. All draft KV writes and collected hidden rows remain +present. Dynamic MTP, decode overlap, and other models retain the existing draft +layout. Expanded graphs consume additional graph memory; capacity settings are +unchanged, but this is not a claim of identical total GPU memory usage. + +See [H200 measurements](h200_optimizations.md) for paired +throughput, full GSM8K results, and the numerical tradeoff. Performance depends on +batch size and input/output length; the main-scene gain is not a universal bound. + ## State lifecycle and precision Verify overwrites the uncommitted history suffix. Acceptance advances only the @@ -68,6 +95,9 @@ clears the tail at aligned pool boundaries. ## Memory and performance +See [H200 / Qwen3.5-27B optimization measurements](h200_optimizations.md) +for workload coverage, numerical tradeoffs and benchmark commands. + For a TP4 model with 48 GDN layers, 12 local value heads and 128x128 BF16 state, one full-model checkpoint occupies 18 MiB per GPU. Native MTP3 uses 72 MiB per request. Replay L8/precompute uses about 24.07 MiB including history and cursors. diff --git a/docs/replayssm/h200_optimizations.md b/docs/replayssm/h200_optimizations.md new file mode 100644 index 0000000000..ac034ee911 --- /dev/null +++ b/docs/replayssm/h200_optimizations.md @@ -0,0 +1,125 @@ +# H200 / Qwen3.5-27B：负载覆盖与正常入口验证 + +2026-10-13。固定原 PR `3af33c2f47f1738e17eb70a3c1e8f42e25ed65ed`、TP1、MTP4、BF16 SSM、L16、264 请求容量和 65536 KV token。所有输入加输出不超过 4096 token。没有通过增加资源容量或更改 MTP 配置获得收益。 + +本次包含 GPU2 的 12 组负载、72 轮交替配对（144 个 cohort），以及 GPU5 正常服务入口的独立整合验证。正常入口由 `python -m lightllm.server.api_server` 启动；候选额外启用 `--replayssm_projection_mode parallel --replayssm_fold_method derived`。 + +## 负载矩阵：同卡交替配对 + +每组、每臂预热两轮，再进行六轮正反顺序交替。greedy、关闭 thinking、禁用 prompt cache 查询、固定输出长度并 ignore_eos。相同 rep 使用相同提示。短输入使用四类技术解释模板及变化的示例编号；长输入使用固定 filler，不能代表所有生产内容。吞吐包含 prefill、请求提交和尾部完成，按总输出 token / 总 wall time 计算。95%区间为10000次配对 bootstrap，seed=417；没有删掉较慢轮次。 + +| 输入 / 输出 token | 并发 | 原 PR token/s | 候选 token/s | 增幅 | 95%区间 | +| --- | ---: | ---: | ---: | ---: | --- | +| 短输入 / 384 | 1 | 140.8 | 170.4 | 21.01% | [18.85%, 23.11%] | +| 短输入 / 384 | 8 | 961.6 | 1059.4 | 10.17% | [5.97%, 14.51%] | +| 短输入 / 384 | 32 | 2551.0 | 2956.0 | 15.88% | [13.77%, 18.30%] | +| 短输入 / 384 | 52 | 2680.6 | 3144.7 | 17.31% | [15.00%, 19.74%] | +| 短输入 / 384 | 53 | 2437.2 | 3177.4 | 30.37% | [27.69%, 33.46%] | +| 短输入 / 384 | 64 | 3083.4 | 3669.0 | 18.99% | [17.51%, 20.16%] | +| 短输入 / 384 | 128 | 3655.0 | 4467.6 | 22.23% | [20.58%, 23.66%] | +| 短输入 / 128 | 32 | 2218.9 | 2465.0 | 11.09% | [5.53%, 14.76%] | +| 短输入 / 128 | 64 | 2585.1 | 2991.6 | 15.72% | [8.73%, 22.32%] | +| 短输入 / 128 | 128 | 3077.0 | 3639.8 | 18.29% | [13.34%, 22.48%] | +| 1024 / 384 | 32 | 1558.6 | 1721.5 | 10.46% | [8.53%, 12.00%] | +| 3072 / 384 | 16 | 705.8 | 741.3 | 5.04% | [3.40%, 7.16%] | + +B128 / 输出384的主场景再次超过20%；其他负载不能套用这一数字。B52→B53跨过首轮draft的原图覆盖边界，组合收益从17.3%变为30.4%，与图扩展/accepted-tail投影机制一致，但该比较不是只切换一个组件的消融。输出缩短或长输入占比增加，端到端收益降低。 + +## 延迟 + +| 输入 / 输出 | 并发 | TTFT p50 原PR→候选 ms | TPOT p50 原PR→候选 ms | +| --- | ---: | ---: | ---: | +| 短 / 384 | 1 | 142.8 → 141.4 | 6.74 → 5.52 | +| 短 / 384 | 8 | 219.8 → 163.9 | 7.24 → 6.71 | +| 短 / 384 | 32 | 370.6 → 383.6 | 10.47 → 8.88 | +| 短 / 384 | 52 | 590.8 → 570.8 | 16.85 → 14.41 | +| 短 / 384 | 53 | 600.7 → 534.1 | 19.20 → 14.51 | +| 短 / 384 | 64 | 633.9 → 602.2 | 18.13 → 15.01 | +| 短 / 384 | 128 | 974.1 → 944.8 | 31.04 → 25.06 | +| 短 / 128 | 32 | 359.9 → 367.5 | 10.64 → 9.12 | +| 短 / 128 | 64 | 594.8 → 600.4 | 18.77 → 15.53 | +| 短 / 128 | 128 | 946.6 → 922.3 | 32.73 → 26.42 | +| 1024 / 384 | 32 | 2185.1 → 2218.5 | 13.15 → 11.31 | +| 3072 / 384 | 16 | 4070.9 → 4021.2 | 11.19 → 10.32 | + +延迟是客户端观测值，按同一负载的请求合并取分位数；TPOT为每请求首尾token时间差除以后续token数。TTFT没有一致改善。逐轮数据保留在本机 `/tmp/replayssm-h200/coverage/`。 + +## 正常服务入口 + +正常入口两臂均预热两轮并正式运行八轮。GPU5原 PR **3667.2 token/s** → 工作树候选 **4446.3 token/s**，增幅 **21.25%**，独立bootstrap95%区间 **[18.96%, 23.66%]**。候选完成后停止服务，再在同一GPU5加载原PR；这是顺序独立运行，不是交替配对。区间反映轮次波动，不覆盖两次运行之间的系统漂移。GPU2的交替配对用于补充因果证据，不能混为同一个统计设计。 +正常入口TTFT p50 原PR 913.2 → 候选 1031.2 ms；TPOT 31.21 → 24.88 ms。正常入口正式性能阶段没有其他模型客户端并发压测。 + +候选已经通过B1/B52/B53/B128的正常请求 smoke。HTTP 随机采样、top-k/top-p、非默认presence/frequency/repetition惩罚、EOS最短长度与无效token参数请求全部完成，logprob有限。Prefix checkpoint 两次实际命中1536 token，输出64 token与禁用查询的请求逐项一致。原始脚本中的 `cache_cold` 标签并非真正冷缓存：禁用查询的首次请求仍填充了缓存，以实际命中长度为准。 + +## 质量与数值权衡 + +同样的全量GSM8K1319题，greedy、关闭thinking、并发64、最大输出1536：正常入口候选 **1278/1319**，历史原 PR **1279/1319**；截断8/6。原PR独对3题、候选独对2题，879条完整token序列一致。准确率差 −0.0758个百分点，配对bootstrap95%区间 [−0.3791,+0.2274] 个百分点。这是与保存的历史质量运行比较，不是本次同卡交替质量实验，不能证明严格无损。此前实验入口 derived 为1281/1319；正常入口和实验入口1162条完整序列一致。 + +对4道有变化的题另做小并发诊断，1536上限1/4正确、截断2；3072上限2/4正确、截断1。原PR在同样小并发复测中，1536上限2/4正确、截断1，3072上限3/4正确、无截断。该诊断改变了并发，且是事后选择的困难题，只说明部分截断可以缓解；没有替换全量测试分数，也不能将差异全部归因于输出上限。 + +并行求解采用activation精度的Tensor Core投影；derived fold复用现有FP32 keys/deltas、以FP32矩阵累加合并历史。历史容量和精度未降低，但运算顺序与checkpoint舍入变化，不能保证与原PR逐位一致。默认仍为raw fold，derived显式启用；固定共享输入的长期测试覆盖 L16/8192 步及 L64/4096 步；这不等同于模型反馈下的质量保证。 + +## 保留的实现与清理范围 + +保留 parallel 求解、可选 derived fold、首轮 draft 图扩展与 accepted-tail 投影、LM head 布局、strided gate、融合 greedy、draft argmax 和默认 token 惩罚跳过。单独的 parallel 求解已测得约 11.4% 端到端提升；其余组件的组合进一步改善吞吐，不能把组合收益算作每个组件的独立收益。 + +删除批量 checkpoint、普通 decode prepare/accept 融合，以及早期特定形状的 H200 调参文件：前两项没有可靠的端到端收益，后者不用于已验证的 parallel 固定布局。恢复普通路径的 autotune 版本，保留 parallel/fold 配置隔离。旧实验原型、重复源码快照、日志、图表和过时报告移出仓库。 + +以上性能数字记录的是清理前的完整候选，不能冒充清理后的重新测量。删除的普通 decode 融合不用于 MTP4 verify；批量 checkpoint 对已测场景收益可忽略。清理后的验证针对状态、checkpoint、CUDA Graph、MTP 和采样回归。 + +## 清理后的验证 + +本地 H200 回归覆盖 676 个用例：首轮 652 通过，24 个 autotune 图隔离参数组合因旧测试桩缺少 `logits_row_stride` 而失败；补齐该字段后，autotune 与最终版采样的 74 项复测全部通过，没有未解决的失败。覆盖普通及 MTP 状态生命周期、checkpoint、GDN/KDA、CUDA Graph、parallel verify width 2/3/5/8/16、derived fold、模型输出裁剪和采样边界。计数包含复测重叠，不相加。 + +24 个修改/新增 Python 文件通过语法检查和仓库 flake8 规则，`git diff --check` 通过。10 月 13 日清理时未重跑端到端模型吞吐或 GSM8K；上述性能和质量数字均来自清理前实验。清理后的吞吐复测见下一节。日志位于本机 `/tmp/replayssm-h200/cleanup-regression.log` 与 `cleanup-recheck.log`。 + +## 2026-10-14：清理后吞吐复测 + +同一 H200 GPU5、Qwen3.5-27B、TP1、MTP4、BF16 SSM、L16、264 请求容量、65536 KV token。并发 128，每请求固定生成 384 token，沿用上一轮短输入技术解释模板，greedy、关闭 thinking、禁用 prompt cache 查询、ignore_eos。两边都通过正常 API 入口运行，各预热 2 轮、正式测量 12 轮。原 PR 仍为 `3af33c2f47f1738e17eb70a3c1e8f42e25ed65ed`；候选为清理后的代码快照，显式启用 `parallel + derived`。 + +| 指标 | 原 PR | 清理后候选 | +| --- | ---: | ---: | +| 端到端吞吐，token/s | 3663.12 | 4505.88 | +| TTFT p50，ms | 928.97 | 901.96 | +| TPOT p50，ms | 31.10 | 24.91 | +| 正式请求数 | 1536 | 1536 | +| 输出 token 数 | 589824 | 589824 | +| prompt cache 命中请求 | 0 | 0 | +| 性能窗最高温度，°C | 63 | 62 | +| 性能窗显存峰值，MiB | 104887 | 104139 | + +**吞吐提升 23.01%，独立 bootstrap 95% 区间 [21.26%, 24.52%]**（10000 次重采样，seed=417）。按全部输出 token / 全部 cohort wall time 汇总，不删除慢轮次。原 PR 的单轮范围为 3572–3741 token/s，候选为 4285–4597 token/s。全部 1179648 个输出 token 的 logprob 有限，两边都没有 HW/SW 热降频。每 verify step 输出 token 数为 3.1833 / 3.1773，候选未通过提高接受率获得这一收益。 + +这是先基线、后候选的同卡顺序独立运行，**不是交替配对**。区间反映轮次波动，不覆盖两个运行阶段之间的系统漂移；本次 23.01% 与上轮 21.25% 的差异不能归因于代码清理本身。本次仅复测这一主场景，没有重新运行其他负载或 GSM8K。结果表明移除低收益路径后，该场景的吞吐提升仍保留。 + +两份源码均已冻结并保存 Python 文件哈希；候选快照与测试结束时的工作区一致。启动参数除代码路径、优化开关和端口外一致。候选首次因旧 NCCL 端口占用而未启动成功，改用空闲端口后完成全部测量；失败启动没有产生性能样本。测试服务均已停止。 + +原始逐轮响应、客户端、启动命令、源码哈希、硬件记录和统计脚本保存在本机 `/tmp/replayssm-h200/cleanup_retest_20261014/`；[紧凑统计结果](h200_retest_summary.json)保留逐轮吞吐与汇总值。 + +## 复现与原始记录 + +正常入口使用相同模型、MTP 配置和容量，候选额外加 `--replayssm_projection_mode parallel --replayssm_fold_method derived`： + +```bash +LIGHTLLM_TRITON_AUTOTUNE_LEVEL=3 LOADWORKER=4 python -m lightllm.server.api_server \ + --model_dir /models/Qwen3.5-27B --tp 1 --host 127.0.0.1 --port 19131 --nccl_port 19132 \ + --disable_vision --disable_audio --ssm_state_mode replay --replayssm_cache_len 16 \ + --linear_att_ssm_data_type bfloat16 --llm_prefill_att_backend triton triton \ + --llm_decode_att_backend fa3 triton --graph_max_batch_size 1280 \ + --running_max_req_size 264 --max_req_total_len 4096 --max_total_token_num 65536 \ + --batch_max_tokens 4096 --chunked_prefill_size 512 --linear_att_hash_page_size 512 \ + --linear_att_page_block_num 2 --linear_att_cache_size 16 --mtp_mode eagle_with_att --mtp_step 4 \ + --replayssm_projection_mode parallel --replayssm_fold_method derived +``` + +通用 HTTP benchmark 保留在 `test/benchmark/kernels/benchmark_replayssm_service.py`，用于固定 1024-token 输入的跨版本对照；它的提示与上述短输入实验不同，不能直接复现同一吞吐数字。两臂必须使用相同参数和预热次数。 + +```bash +python test/benchmark/kernels/benchmark_replayssm_service.py \ + --model-dir /models/Qwen3.5-27B --url http://127.0.0.1:19131/generate_stream \ + --batch 32 --tokens 384 --warmups 2 --repeats 8 --output /tmp/replayssm-service.json +``` + +原始实验客户端、逐轮结果和启动命令在本机 `/tmp/replayssm-h200/{coverage,integrated,normal_baseline}/`。清理前代码及文档完整归档于 `/tmp/replayssm-h200/cleanup-before.tar.gz`。这些是本地实验记录，不是仓库依赖。 + +GPU2 矩阵性能窗最高 60°C，正常入口候选/原 PR 的 GPU5 最高 62/63°C；均无硬件或软件热降频。GPU2 部分测试与其他 GPU 的模型加载/质量测试共享主机。正常入口性能阶段没有其他模型客户端并发压测。正常入口显存峰值候选 104921 MiB、原 PR 104785 MiB；扩展图使候选多占约 136 MiB，容量配置相同不代表总显存完全相等。 diff --git a/docs/replayssm/h200_retest_summary.json b/docs/replayssm/h200_retest_summary.json new file mode 100644 index 0000000000..f205d2024e --- /dev/null +++ b/docs/replayssm/h200_retest_summary.json @@ -0,0 +1,86 @@ +{ + "design": "GPU5, normal API entry, baseline then candidate sequential independent runs; 2 warmups and 12 measured cohorts per arm", + "baseline_commit": "3af33c2f47f1738e17eb70a3c1e8f42e25ed65ed", + "batch": 128, + "output_tokens": 384, + "arms": { + "baseline": { + "repeats": 12, + "requests": 1536, + "output_tokens": 589824, + "tokens_per_s": 3663.117912349209, + "cohort_tokens_per_s": [ + 3572.39450013313, + 3679.1945250855033, + 3668.5673855693176, + 3664.162821283357, + 3655.0134290105893, + 3658.030321178895, + 3741.2360562287436, + 3733.451128039074, + 3630.357737495093, + 3637.8124518110835, + 3689.765095342039, + 3633.695617211311 + ], + "finite_logprobs": true, + "ttft_p50_ms": 928.9662055671215, + "tpot_p50_ms": 31.098604479478794, + "hardware": { + "samples": 139, + "max_temperature_C": 63.0, + "max_memory_MiB": 104887.0, + "min_sm_clock_MHz": 1665.0, + "max_sm_clock_MHz": 1980.0, + "thermal_slowdown_samples": 0 + }, + "acceptance": { + "logged_requests": 1536, + "prompt_cache_hits": 0, + "output_per_verify": 3.1832994219777966 + } + }, + "candidate": { + "repeats": 12, + "requests": 1536, + "output_tokens": 589824, + "tokens_per_s": 4505.876589509283, + "cohort_tokens_per_s": [ + 4401.733921373901, + 4558.013559046327, + 4454.915305507067, + 4543.451463667397, + 4460.422265583761, + 4284.845559058763, + 4583.72309215793, + 4597.109123661394, + 4551.504440860731, + 4579.715485515799, + 4497.592803040789, + 4579.283607210278 + ], + "finite_logprobs": true, + "ttft_p50_ms": 901.9625158980489, + "tpot_p50_ms": 24.913206290424494, + "hardware": { + "samples": 112, + "max_temperature_C": 62.0, + "max_memory_MiB": 104139.0, + "min_sm_clock_MHz": 1500.0, + "max_sm_clock_MHz": 1980.0, + "thermal_slowdown_samples": 0 + }, + "acceptance": { + "logged_requests": 1536, + "prompt_cache_hits": 0, + "output_per_verify": 3.17729762924417 + } + } + }, + "gain_pct": 23.006594309152373, + "independent_bootstrap_ci95_pct": [ + 21.26323367507323, + 24.51623928708578 + ], + "caveat": "The interval measures round-to-round variation, not drift between sequential service runs. This benchmark does not re-evaluate model quality." +} diff --git a/lightllm/common/basemodel/basemodel.py b/lightllm/common/basemodel/basemodel.py index 02478d3e0b..22c2cc4684 100755 --- a/lightllm/common/basemodel/basemodel.py +++ b/lightllm/common/basemodel/basemodel.py @@ -276,6 +276,21 @@ def _align_decode_batch_size(self, batch_size: int) -> int: alignment = math.lcm(alignment, self.tp_world_size_) return triton.cdiv(batch_size, alignment) * alignment + def use_replayssm_draft_tail(self): + """Keep first-draft graph shapes distinct from single-row draft decode.""" + return ( + self.is_mtp_draft_model + and self.args.ssm_state_mode == "replay" + and self.args.replayssm_projection_mode == "parallel" + and self.args.mtp_mode == "eagle_with_att" + and self.args.mtp_step > 0 + and not self.args.mtp_dynamic_verify + and not self.args.enable_decode_microbatch_overlap + and not self.args.enable_tpsp_mix_mode + and self.tp_world_size_ == 1 + and self.config.get("model_type") in ("qwen3_5", "qwen3_5_text") + ) + def _init_cudagraph(self): # When graph covers the configured request length, it must also cover MTP's internal token margin. if self.args.mtp_mode is not None and self.graph_max_len_in_batch >= self.args.max_req_total_len: @@ -297,6 +312,11 @@ def _init_cudagraph(self): ) ) if self.graph is not None: + if self.use_replayssm_draft_tail(): + width = self.args.mtp_step + 1 + sizes = self.graph.cuda_graph_batch_sizes + self.graph.cuda_graph_batch_sizes = sorted(set(sizes + [width * n for n in sizes])) + self.graph.max_batch_size = self.graph.cuda_graph_batch_sizes[-1] if get_env_start_args().enable_decode_microbatch_overlap: self.graph.warmup_overlap(self) else: @@ -354,6 +374,24 @@ def _create_inferstate(self, model_input: ModelInput, microbatch_index: int = 0) infer_state.is_draft_model = self.is_mtp_draft_model infer_state.return_all_prompt_logics = self.return_all_prompt_logics infer_state.batch_size = model_input.batch_size + # Above the logical request capacity, this can only be the first draft + # pass. Warmup and replay must use the same compact projection layout. + if ( + self.use_replayssm_draft_tail() + and not model_input.is_prefill + and model_input.batch_size > self.args.running_max_req_size + ): + width = self.args.mtp_step + 1 + assert model_input.batch_size % width == 0 + count = model_input.batch_size // width + rows = model_input.b_logits_rows + if rows is None: + rows = torch.arange(count, device=model_input.input_ids.device, dtype=torch.int64) * width + width - 1 + else: + assert rows.ndim == 1 and rows.numel() <= count + rows = F.pad(rows, (0, count - rows.numel())) + infer_state.b_logits_rows = rows + infer_state.logits_row_stride = width infer_state.total_token_num = model_input.total_token_num infer_state.max_q_seq_len = model_input.max_q_seq_len infer_state.max_kv_seq_len = model_input.max_kv_seq_len @@ -492,17 +530,21 @@ def _create_model_output(self, post_output: PostLayerOutput, infer_state: InferS prompt_logics=infer_state.prompt_logics, logits_token_ids=post_output.logits_token_ids, ssm_run_config=infer_state.ssm_run_config, + logits_row_stride=infer_state.logits_row_stride, ) return output def _create_unpad_decode_model_output(self, model_output: ModelOutput, origin_batch_size: int): - padded_batch_size = model_output.logits.shape[0] + stride = model_output.logits_row_stride + padded_batch_size = model_output.logits.shape[0] * stride + assert origin_batch_size % stride == 0 + logits_size = origin_batch_size // stride if padded_batch_size == origin_batch_size: return model_output new_model_output = copy.copy(model_output) - new_model_output.logits = new_model_output.logits[0:origin_batch_size] + new_model_output.logits = new_model_output.logits[0:logits_size] if new_model_output.logits_token_ids is not None: - new_model_output.logits_token_ids = new_model_output.logits_token_ids[:origin_batch_size] + new_model_output.logits_token_ids = new_model_output.logits_token_ids[:logits_size] new_model_output.mtp_collector = model_output.mtp_collector.unpad_decode( padded_batch_size=padded_batch_size, origin_batch_size=origin_batch_size, @@ -631,7 +673,6 @@ def _decode( @final def _context_forward(self, infer_state: InferStateInfo): - input_embs = self.pre_infer.context_forward(infer_state.input_ids, infer_state, self.pre_post_weight) if self.args.enable_dp_prefill_balance: assert not self.args.enable_prefill_cudagraph, "not support now" diff --git a/lightllm/common/basemodel/batch_objs.py b/lightllm/common/basemodel/batch_objs.py index a393c5c3e0..f73c2dc153 100644 --- a/lightllm/common/basemodel/batch_objs.py +++ b/lightllm/common/basemodel/batch_objs.py @@ -52,8 +52,13 @@ class ModelInput: # 的 draft 模型的输入 mtp_draft_input_hiddens: Optional[torch.Tensor] = None + # Optional accepted-tail rows for the first EAGLE draft projection. + b_logits_rows: Optional[torch.Tensor] = None + def to_cuda(self): self.check_input() + if self.b_logits_rows is not None: + self.b_logits_rows = self.b_logits_rows.cuda(non_blocking=True) # Prefill 和 decode 都必须提供的公共张量。 self.b_req_idx = self.b_req_idx.cuda(non_blocking=True) @@ -204,6 +209,8 @@ class ModelOutput: # Layout selected by this verify graph; preserved when logits are unpadded. ssm_run_config: Optional[dict] = None + # Number of physical decode rows represented by each selected logits row. + logits_row_stride: int = 1 def __post_init__(self) -> None: if self.mtp_collector is None: diff --git a/lightllm/common/basemodel/infer_struct.py b/lightllm/common/basemodel/infer_struct.py index 26ef020348..a88c629a81 100755 --- a/lightllm/common/basemodel/infer_struct.py +++ b/lightllm/common/basemodel/infer_struct.py @@ -65,6 +65,8 @@ def __init__(self): # 在一些细节场景下需要有该信息区分一些资源的申请和管理。 self.microbatch_index: int = 0 self.ssm_run_config: Optional[dict] = None + self.b_logits_rows: Optional[torch.Tensor] = None + self.logits_row_stride: int = 1 # 当前 forward 独占的 hidden state 收集器。普通推理使用短生命周期实例； # Prefill CUDA Graph 使用随 graph infer state 长期保存的实例。 diff --git a/lightllm/common/basemodel/triton_kernel/linear_att/replayssm.py b/lightllm/common/basemodel/triton_kernel/linear_att/replayssm.py index 82eb65e67f..17c3a1cfc5 100644 --- a/lightllm/common/basemodel/triton_kernel/linear_att/replayssm.py +++ b/lightllm/common/basemodel/triton_kernel/linear_att/replayssm.py @@ -48,6 +48,7 @@ def _prepare_fold( VARLEN: tl.constexpr, BLOCK: tl.constexpr, WIDTH: tl.constexpr, + SEPARATE_FOLD: tl.constexpr, ): seq = tl.arange(0, BLOCK) req = tl.load(Reqs + seq, seq < N, HOLD) @@ -71,6 +72,8 @@ def _prepare_fold( if N >= 256 and L == 8 and WIDTH == 4: short = tl.sum((valid & (length > 0) & (length < WIDTH)).to(tl.int32), axis=0) split = split & ((live < 256) | (short > 0)) + if SEPARATE_FOLD: + split = count > 0 tl.store(ForwardCursors + req, tl.where(overflow & split, (cursor & (2 * L)) ^ (2 * L), cursor), valid) tl.store(Active + N, tl.where(split, count, 0)) @@ -115,6 +118,9 @@ def _fold_history( GS: tl.constexpr, BS: tl.constexpr, KDA: tl.constexpr, + Keys, + Deltas, + DERIVED: tl.constexpr, ): count = tl.load(Active + N) work = tl.program_id(0) @@ -136,15 +142,36 @@ def _fold_history( key_slot = req * KH + hv // (HV // KH) sp = State + layer * SS + 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 + layer * RS + (key_slot * (2 * L) + base + j) * K + kk, kk < K, 0).to(tl.float32) - raw_v = tl.load(RawValues + layer * VS + (slot * (2 * L) + base + j) * V + vv, vv < V, 0).to(tl.float32) - if KDA: - g = tl.load(Gates + layer * GS + (slot * (2 * L) + base + j) * K + kk, kk < K, 0).to(tl.float32) - else: - g = tl.load(Gates + layer * GS + slot * (2 * L) + base + j).to(tl.float32) - beta = tl.load(Betas + layer * BS + slot * (2 * L) + base + j).to(tl.float32) - state = _replay_update(state, raw_k, raw_v, g, beta, KDA) + if DERIVED: + tl.static_assert(not KDA) + ll = tl.arange(0, max(16, L)) + g = tl.load(Gates + layer * GS + slot * (2 * L) + base + ll, ll < n, 0).to(tl.float32) + prefix = tl.cumsum(g, 0) + total_g = tl.sum(g, 0) + keys = tl.load( + Keys + layer * RS + (key_slot * (2 * L) + base + ll[None, :]) * K + kk[:, None], + (kk[:, None] < K) & (ll[None, :] < n), + 0, + ).to(tl.float32) + delta = tl.load( + Deltas + layer * VS + (slot * (2 * L) + base + ll[:, None]) * V + vv[None, :], + (ll[:, None] < n) & (vv[None, :] < V), + 0, + ).to(tl.float32) + keys *= tl.exp(total_g - prefix)[None, :] + state = state * tl.exp(total_g) + tl.dot(keys, delta, input_precision="tf32x3") + else: + for j in range(n): + raw_k = tl.load(RawKeys + layer * RS + (key_slot * (2 * L) + base + j) * K + kk, kk < K, 0).to( + tl.float32 + ) + raw_v = tl.load(RawValues + layer * VS + (slot * (2 * L) + base + j) * V + vv, vv < V, 0).to(tl.float32) + if KDA: + g = tl.load(Gates + layer * GS + (slot * (2 * L) + base + j) * K + kk, kk < K, 0).to(tl.float32) + else: + g = tl.load(Gates + layer * GS + slot * (2 * L) + base + j).to(tl.float32) + beta = tl.load(Betas + layer * BS + slot * (2 * L) + base + j).to(tl.float32) + state = _replay_update(state, raw_k, raw_v, g, beta, KDA) # Start the new history from the same rounded checkpoint future calls load. state = state.to(State.dtype.element_ty).to(tl.float32) tl.store(sp, state, (kk[:, None] < K) & (vv[None, :] < V)) @@ -191,6 +218,7 @@ def _replay( FUSED_FOLD: tl.constexpr, KDA: tl.constexpr, LOWER: tl.constexpr, + PARALLEL_SOLVE: tl.constexpr, ): """Output reconstruction with an exact raw-input recurrence at checkpoint folds.""" iv, seq, hv = tl.program_id(0), tl.program_id(1), tl.program_id(2) @@ -252,8 +280,11 @@ def _replay( 0, ).to(tl.float32) qk = q_rows + k_rows - qk *= tl.rsqrt(tl.sum(qk * qk, axis=1) + 1.0e-6)[:, None] - qk *= tl.where(tt < BW, K ** -0.5, 1.0)[:, None] + if PARALLEL_SOLVE: + qk /= tl.sqrt(tl.sum(qk * qk, axis=1) + 1.0e-6)[:, None] + else: + qk *= tl.rsqrt(tl.sum(qk * qk, axis=1) + 1.0e-6)[:, None] + qk *= tl.where(tt < BW, K**-0.5, 1.0)[:, None] if KDA: history = ll < n gates = tl.load( @@ -331,14 +362,44 @@ def _replay( 0, ).to(tl.float32) betas = tl.sigmoid(tl.load(B + fresh_tokens * SB + hv, fresh, 0).to(tl.float32)) - for t in range(end - start): - checkpoint_sk = tl.sum(tl.where((tt == BW + t)[:, None], projections, 0), 0) - weights = tl.sum(tl.where((tt == BW + t)[:, None], coefficients, 0), 0) - sk = checkpoint_sk + tl.sum(deltas * weights[:, None], 0) - value = tl.sum(tl.where((jj == L + t)[:, None], values, 0), 0) - beta = tl.sum(tl.where(jj == L + t, betas, 0)) - delta = beta * (value - sk) - deltas = tl.where((jj == L + t)[:, None], delta[None, :], deltas) + if PARALLEL_SOLVE: + projected_history = projections + tl.dot( + coefficients.to(projection_dtype), deltas.to(projection_dtype), input_precision=projection_precision + ) + wi = tl.arange(0, max(16, 2 * BW)) + active = (wi < BW) & (start + wi < end) + key_rows = tl.gather( + projected_history, tl.broadcast_to((BW + wi % BW)[:, None], projected_history.shape), 0 + ) + fresh_values = tl.gather(values, tl.broadcast_to((L + wi % BW)[:, None], (max(16, 2 * BW), BV)), 0) + beta_rows = tl.sigmoid(tl.load(B + (start + wi) * SB + hv, active, 0).to(tl.float32)) + rhs = tl.where(active[:, None], beta_rows[:, None] * (fresh_values - key_rows), 0.0) + cross = tl.gather(coefficients, tl.broadcast_to((BW + wi % BW)[:, None], coefficients.shape), 0) + cross = tl.gather(cross, tl.broadcast_to((L + wi % BW)[None, :], (max(16, 2 * BW), max(16, 2 * BW))), 1) + lower = active[:, None] & active[None, :] & (wi[:, None] > wi[None, :]) + power = tl.where(lower, -beta_rows[:, None] * cross, 0.0) + inverse = (wi[:, None] == wi[None, :]).to(tl.float32) + # The strictly lower-triangular coupling is nilpotent. Repeated + # squaring builds (I - power)^-1 in log2(BW) matrix steps; capacity + # is at most 64, so six steps cover every supported verify width. + for solve in tl.static_range(6): + if (1 << solve) < BW: + inverse += tl.dot(power, inverse, input_precision="tf32x3") + power = tl.dot(power, power, input_precision="tf32x3") + solved = tl.dot( + inverse.to(projection_dtype), rhs.to(projection_dtype), input_precision=projection_precision + ) + mapped = tl.gather(solved, tl.broadcast_to(((jj - L) % BW)[:, None], deltas.shape), 0) + deltas = tl.where(fresh[:, None], mapped, deltas) + else: + for t in range(end - start): + checkpoint_sk = tl.sum(tl.where((tt == BW + t)[:, None], projections, 0), 0) + weights = tl.sum(tl.where((tt == BW + t)[:, None], coefficients, 0), 0) + sk = checkpoint_sk + tl.sum(deltas * weights[:, None], 0) + value = tl.sum(tl.where((jj == L + t)[:, None], values, 0), 0) + beta = tl.sum(tl.where(jj == L + t, betas, 0)) + delta = beta * (value - sk) + deltas = tl.where((jj == L + t)[:, None], delta[None, :], deltas) outputs = projections + tl.dot( coefficients.to(projection_dtype), deltas.to(projection_dtype), input_precision=projection_precision ) @@ -406,7 +467,7 @@ def _replay( 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) + 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) if KDA: x = tl.load(A + t * SA + hv * K + kk, kk < K, 0).to(tl.float32) + bias @@ -532,10 +593,19 @@ def __init__( lower_bound=-5.0, projection_mode="inline", run_config=None, + fold_method="raw", ): assert state.dtype in (torch.float32, torch.bfloat16), "ReplaySSM requires FP32 or BF16 SSM state" assert verify_width > 0 and capacity >= max(4, verify_width) and capacity & (capacity - 1) == 0 - assert projection_mode in ("inline", "precompute") + assert projection_mode in ("inline", "precompute", "parallel") + assert fold_method in ("raw", "derived") + assert fold_method == "raw" or projection_mode == "parallel" + if projection_mode == "parallel": + assert not kda and verify_width > 1, "Parallel ReplaySSM requires GDN MTP" + assert state.shape[-2:] == (128, 128), "Parallel ReplaySSM requires 128x128 heads" + assert activation_dtype == torch.bfloat16, "Parallel ReplaySSM requires BF16 activations" + assert state.dtype == torch.bfloat16, "Parallel ReplaySSM currently supports BF16 state" + self.fold_method = fold_method self.state = state self.kda = kda self.lower_bound = lower_bound @@ -600,6 +670,7 @@ def prepare_decode(self, reqs, cu_seqlens=None): cu_seqlens is not None, triton.next_power_of_2(reqs.numel()), self.verify_width, + self.projection_mode == "parallel", num_warps=4, ) programs = min(self.fold_programs, triton.cdiv(vd, 32) * reqs.numel() * layers * hv) @@ -627,6 +698,9 @@ def prepare_decode(self, reqs, cu_seqlens=None): self.gates.stride(0), self.betas.stride(0), self.kda, + self.keys, + self.deltas, + self.fold_method == "derived", num_warps=2, num_stages=1, ) @@ -688,6 +762,8 @@ def snapshot_accepted_state(self, req_idx): return output def get_run_config(self, q, k, v, a, b, a_log, bias, cu_seqlens=None): + if self.projection_mode == "parallel": + return self.run_config or {"BV": 128, "num_warps": 2, "num_stages": 1} return configure_cache(self, "replay", q, k, v, a, b, a_log, bias, cu_seqlens) def forward(self, layer, q, k, v, a, b, a_log, bias, reqs, cu_seqlens=None, run_config=None): @@ -697,7 +773,9 @@ def forward(self, layer, q, k, v, a, b, a_log, bias, reqs, cu_seqlens=None, run_ axis = 1 if cu_seqlens is not None else 0 out = torch.empty_like(v) bv = config["BV"] - precompute_state = self.verify_width > 1 and bv >= 16 and self.projection_mode == "precompute" + precompute_state = self.verify_width > 1 and bv >= 16 and self.projection_mode in ("precompute", "parallel") + if self.projection_mode == "parallel": + assert precompute_state, "Parallel ReplaySSM requires BV >= 16" _replay[(triton.cdiv(vd, bv), reqs.numel(), hv)]( q, k, @@ -734,9 +812,10 @@ def forward(self, layer, q, k, v, a, b, a_log, bias, reqs, cu_seqlens=None, run_ cu_seqlens is not None, precompute_state, triton.next_power_of_2(self.verify_width), - reqs.numel() > 12, + reqs.numel() > 12 and self.projection_mode != "parallel", self.kda, self.lower_bound, + self.projection_mode == "parallel", num_warps=config["num_warps"], num_stages=config.get("num_stages", 3), ) diff --git a/lightllm/common/basemodel/triton_kernel/linear_att/ssm_autotune.py b/lightllm/common/basemodel/triton_kernel/linear_att/ssm_autotune.py index f3d5ab06ac..681935f953 100644 --- a/lightllm/common/basemodel/triton_kernel/linear_att/ssm_autotune.py +++ b/lightllm/common/basemodel/triton_kernel/linear_att/ssm_autotune.py @@ -37,6 +37,10 @@ def static_key(cache, mode, q, k, v, a, b, cu_seqlens): key["lower"] = cache.lower_bound if mode == "replay": key["history_dtype"] = str(cache.keys.dtype) + if cache.projection_mode == "parallel": + key["projection"] = "parallel" + key["fold"] = cache.fold_method + return key key["projection"] = ( ("kda_mma" if cache.kda else "block_mma_activation_precision") if cache.projection_mode == "precompute" and cache.verify_width > 1 @@ -65,6 +69,7 @@ def rebuild_inputs(cache, mode, q, k, v, a, b, a_log, bias, cu_seqlens=None, wor kda=cache.kda, lower_bound=cache.lower_bound, projection_mode=cache.projection_mode, + fold_method=cache.fold_method, ) else: scratch = type(cache)( diff --git a/lightllm/common/basemodel/triton_kernel/post_process/apply_penalty_gpu_cache.py b/lightllm/common/basemodel/triton_kernel/post_process/apply_penalty_gpu_cache.py index a46fd7436c..9d70ea034e 100644 --- a/lightllm/common/basemodel/triton_kernel/post_process/apply_penalty_gpu_cache.py +++ b/lightllm/common/basemodel/triton_kernel/post_process/apply_penalty_gpu_cache.py @@ -89,27 +89,29 @@ def apply_penalty_gpu_cache( b_mask_eos_reqs: torch.Tensor, eos_ids: torch.Tensor, sampling_params_manager: ReqSamplingParamsManager, + skip_token_penalties: bool = False, ): assert Logits.is_contiguous() - BLOCK_P = 2048 - num_warps = 8 - vocab_size = sampling_params_manager.vocab_size - req_to_out_token_id_counter = sampling_params_manager.req_to_out_token_id_counter - _fwd_kernel_apply_penalty_cache[(Logits.shape[0], triton.cdiv(vocab_size, BLOCK_P))]( - Logits=Logits, - stride_logit_b=Logits.stride(0), - stride_logit_s=Logits.stride(1), - b_req_idx=b_req_idx, - req_to_presence_penalty=sampling_params_manager.req_to_presence_penalty, - req_to_frequency_penalty=sampling_params_manager.req_to_frequency_penalty, - req_to_repetition_penalty=sampling_params_manager.req_to_repetition_penalty, - req_to_out_token_id_counter=req_to_out_token_id_counter, - stride_counter_r=req_to_out_token_id_counter.stride(0), - stride_counter_s=req_to_out_token_id_counter.stride(1), - vocab_size=vocab_size, - BLOCK_P=BLOCK_P, - num_warps=num_warps, - ) + if not skip_token_penalties: + BLOCK_P = 2048 + num_warps = 8 + vocab_size = sampling_params_manager.vocab_size + req_to_out_token_id_counter = sampling_params_manager.req_to_out_token_id_counter + _fwd_kernel_apply_penalty_cache[(Logits.shape[0], triton.cdiv(vocab_size, BLOCK_P))]( + Logits=Logits, + stride_logit_b=Logits.stride(0), + stride_logit_s=Logits.stride(1), + b_req_idx=b_req_idx, + req_to_presence_penalty=sampling_params_manager.req_to_presence_penalty, + req_to_frequency_penalty=sampling_params_manager.req_to_frequency_penalty, + req_to_repetition_penalty=sampling_params_manager.req_to_repetition_penalty, + req_to_out_token_id_counter=req_to_out_token_id_counter, + stride_counter_r=req_to_out_token_id_counter.stride(0), + stride_counter_s=req_to_out_token_id_counter.stride(1), + vocab_size=vocab_size, + BLOCK_P=BLOCK_P, + num_warps=num_warps, + ) BLOCK = 128 grid = (triton.cdiv(Logits.shape[0], BLOCK),) diff --git a/lightllm/common/basemodel/triton_kernel/post_process/argmax.py b/lightllm/common/basemodel/triton_kernel/post_process/argmax.py new file mode 100644 index 0000000000..d759a7dd8b --- /dev/null +++ b/lightllm/common/basemodel/triton_kernel/post_process/argmax.py @@ -0,0 +1,47 @@ +"""Two-stage vocabulary argmax with PyTorch tie and NaN semantics.""" + +import torch +import triton +import triton.language as tl + + +@triton.jit +def _argmax_partials(X, M, INDICES, NAN, V: tl.constexpr, S: tl.constexpr, P: tl.constexpr, B: tl.constexpr): + r = tl.program_id(0) + p = tl.program_id(1) + i = p * B + tl.arange(0, B) + x = tl.load(X + r * S + i, i < V, -float("inf")).to(tl.float32) + isnan = x != x + n = tl.min(tl.where((i < V) & isnan, i, 2147483647), 0) + m = tl.max(tl.where(isnan, -float("inf"), x), 0) + idx = tl.min(tl.where((i < V) & ~isnan & (x == m), i, 2147483647), 0) + tl.store(M + r * P + p, m) + tl.store(INDICES + r * P + p, idx) + tl.store(NAN + r * P + p, n) + + +@triton.jit +def _argmax_reduce(M, INDICES, NAN, OUTPUT, P: tl.constexpr, B: tl.constexpr): + r = tl.program_id(0) + i = tl.arange(0, B) + m = tl.load(M + r * P + i, i < P, -float("inf")) + idx = tl.load(INDICES + r * P + i, i < P, 2147483647) + nan = tl.min(tl.load(NAN + r * P + i, i < P, 2147483647), 0) + maximum = tl.max(m, 0) + winner = tl.min(tl.where(m == maximum, idx, 2147483647), 0) + tl.store(OUTPUT + r, tl.where(nan < 2147483647, nan, winner).to(tl.int64)) + + +def argmax(logits): + """Select the first maximum (or first NaN) from contiguous FP32 vocabulary rows.""" + assert logits.dtype == torch.float32 and logits.stride(1) == 1 + rows, vocab = logits.shape + block = 4096 + parts = triton.cdiv(vocab, block) + maxima = torch.empty((rows, parts), device=logits.device, dtype=torch.float32) + indices = torch.empty((rows, parts), device=logits.device, dtype=torch.int32) + nans = torch.empty_like(indices) + output = torch.empty(rows, device=logits.device, dtype=torch.int64) + _argmax_partials[(rows, parts)](logits, maxima, indices, nans, vocab, logits.stride(0), parts, block, num_warps=4) + _argmax_reduce[(rows,)](maxima, indices, nans, output, parts, triton.next_power_of_2(parts), num_warps=4) + return output diff --git a/lightllm/common/basemodel/triton_kernel/post_process/greedy.py b/lightllm/common/basemodel/triton_kernel/post_process/greedy.py new file mode 100644 index 0000000000..827b540163 --- /dev/null +++ b/lightllm/common/basemodel/triton_kernel/post_process/greedy.py @@ -0,0 +1,66 @@ +"""Greedy token selection and log-probability without a full softmax buffer.""" + +import torch +import triton +import triton.language as tl + + +@triton.jit +def _partials(X, T, M, S, INDICES, NAN, V: tl.constexpr, STRIDE: tl.constexpr, P: tl.constexpr, B: tl.constexpr): + row = tl.program_id(0) + part = tl.program_id(1) + i = part * B + tl.arange(0, B) + valid = i < V + x = tl.load(X + row * STRIDE + i, valid, other=-float("inf")).to(tl.float32) + temperature = tl.load(T + row) + x = tl.div_rn(x, temperature) + tl.store(X + row * STRIDE + i, x, valid) + is_nan = x != x + nan = tl.min(tl.where(valid & is_nan, i, 2147483647), 0) + maximum = tl.max(tl.where(is_nan, -float("inf"), x), 0) + winner = tl.min(tl.where(valid & ~is_nan & (x == maximum), i, 2147483647), 0) + denom = tl.sum(tl.where(valid & (maximum != -float("inf")), tl.exp(x - maximum), 0.0), 0) + offset = row * P + part + tl.store(M + offset, maximum) + tl.store(S + offset, denom) + tl.store(INDICES + offset, winner) + tl.store(NAN + offset, nan) + + +@triton.jit +def _finish(M, S, INDICES, NAN, OUT, LOGP, P: tl.constexpr, B: tl.constexpr): + row = tl.program_id(0) + i = tl.arange(0, B) + valid = i < P + m = tl.load(M + row * P + i, valid, other=-float("inf")) + s = tl.load(S + row * P + i, valid, other=0.0) + ids = tl.load(INDICES + row * P + i, valid, other=2147483647) + nans = tl.load(NAN + row * P + i, valid, other=2147483647) + maximum = tl.max(m, 0) + winner = tl.min(tl.where(valid & (m == maximum), ids, 2147483647), 0) + first_nan = tl.min(nans, 0) + winner = tl.where(first_nan < 2147483647, first_nan, winner) + denom = tl.sum(tl.where(valid, s * tl.exp(m - maximum), 0.0), 0) + logp = -tl.log(denom) + logp = tl.where(first_nan < 2147483647, float("nan"), logp) + tl.store(OUT + row, winner.to(tl.int64)) + tl.store(LOGP + row, logp) + + +def greedy_sample(logits, temperatures): + """Scale logits in place, preserving torch.div rounding and first-index argmax ties.""" + assert logits.dtype == torch.float32 and logits.stride(1) == 1 + rows, vocab = logits.shape + block = 4096 + parts = triton.cdiv(vocab, block) + maximum = torch.empty((rows, parts), device=logits.device, dtype=torch.float32) + sums = torch.empty_like(maximum) + indices = torch.empty((rows, parts), device=logits.device, dtype=torch.int32) + nans = torch.empty_like(indices) + ids = torch.empty(rows, device=logits.device, dtype=torch.int64) + logp = torch.empty(rows, device=logits.device, dtype=torch.float32) + _partials[(rows, parts)]( + logits, temperatures, maximum, sums, indices, nans, vocab, logits.stride(0), parts, block, num_warps=4 + ) + _finish[(rows,)](maximum, sums, indices, nans, ids, logp, parts, triton.next_power_of_2(parts), num_warps=4) + return ids, logp diff --git a/lightllm/common/req_manager/linear_att.py b/lightllm/common/req_manager/linear_att.py index 1914421e42..020c64ae9e 100644 --- a/lightllm/common/req_manager/linear_att.py +++ b/lightllm/common/req_manager/linear_att.py @@ -68,6 +68,7 @@ def __init__( num_key_heads=linear_config.num_linear_k_heads, kda=self.kda, projection_mode=args.replayssm_projection_mode, + fold_method=args.replayssm_fold_method, ) elif args.ssm_state_mode == "compact": from lightllm.common.basemodel.triton_kernel.linear_att.replayssm_compact import CompactSSMCache diff --git a/lightllm/models/llama/layer_infer/post_layer_infer.py b/lightllm/models/llama/layer_infer/post_layer_infer.py index 7ca2df0625..c691caa205 100644 --- a/lightllm/models/llama/layer_infer/post_layer_infer.py +++ b/lightllm/models/llama/layer_infer/post_layer_infer.py @@ -102,6 +102,17 @@ def _lm_head_and_gather( ) -> PostLayerOutput: """执行原始 lm-head，并收集完整词表 logits。""" + if self.tp_world_size_ == 1 and hidden.is_cuda and hidden.dtype == torch.bfloat16: + # Produce token-major logits directly so the FP32 conversion reads + # contiguous rows instead of transposing the full vocabulary matrix. + normed = self._norm(hidden, infer_state, layer_weight) + vocab_size = layer_weight.lm_head_weight_.vocab_size + local_logits = self.alloc_tensor((token_num, vocab_size), dtype=hidden.dtype) + torch.mm(normed, layer_weight.lm_head_weight_.weight.t(), out=local_logits) + logits = self.alloc_tensor((token_num, vocab_size), dtype=torch.float32) + logits.copy_(local_logits) + return PostLayerOutput(logits=logits) + local_logits = self._project_local_logits(hidden, token_num, layer_weight, infer_state) vocab_size = layer_weight.lm_head_weight_.vocab_size @@ -230,6 +241,9 @@ def _draft_lm_head_and_gather( ) -> PostLayerOutput: """按 draft 配置返回候选 logits，未开启时回退到完整词表。""" + if infer_state.b_logits_rows is not None: + hidden = hidden.index_select(0, infer_state.b_logits_rows) + token_num = hidden.shape[0] top_k = get_env_start_args().draft_vocab_topk_sampling if top_k is None: return self._lm_head_and_gather(hidden, token_num, layer_weight, infer_state) @@ -240,7 +254,6 @@ def _draft_lm_head_and_gather( def token_forward( self, input_embdings: torch.Tensor, infer_state: LlamaInferStateInfo, layer_weight: LlamaPreAndPostLayerWeight ) -> PostLayerOutput: - return self._token_forward(input_embdings=input_embdings, infer_state=infer_state, layer_weight=layer_weight) def overlap_tpsp_token_forward( @@ -251,7 +264,6 @@ def overlap_tpsp_token_forward( infer_state1: LlamaInferStateInfo, layer_weight: BaseLayerWeight, ) -> tuple[PostLayerOutput, PostLayerOutput]: - output = self.token_forward(input_embdings, infer_state, layer_weight=layer_weight) output1 = self.token_forward(input_embdings1, infer_state1, layer_weight=layer_weight) diff --git a/lightllm/models/qwen3next/layer_infer/transformer_layer_infer.py b/lightllm/models/qwen3next/layer_infer/transformer_layer_infer.py index 92d68c9fd2..4ef9d0b53c 100644 --- a/lightllm/models/qwen3next/layer_infer/transformer_layer_infer.py +++ b/lightllm/models/qwen3next/layer_infer/transformer_layer_infer.py @@ -319,7 +319,7 @@ def _linear_post( ) -> torch.Tensor: num_tokens = z.shape[0] core_attn_out = core_attn_out.view(-1, core_attn_out.shape[-1]) - z = z.contiguous().view(-1, z.shape[-1]) + # Gated RMSNorm accepts the strided [tokens, heads, dim] projection slice. norm_out = layer_weight.linear_norm(core_attn_out, z, self.eps_) core_attn_out = norm_out.view(num_tokens, -1) output = layer_weight.linear_out_proj.mm(core_attn_out) diff --git a/lightllm/server/api_cli.py b/lightllm/server/api_cli.py index b0d970e5c7..b4d90c1ed8 100644 --- a/lightllm/server/api_cli.py +++ b/lightllm/server/api_cli.py @@ -955,9 +955,16 @@ def add_cli_args(parser: argparse.ArgumentParser) -> argparse.ArgumentParser: ) parser.add_argument( "--replayssm_projection_mode", - choices=["inline", "precompute"], + choices=["inline", "precompute", "parallel"], default="inline", - help="ReplaySSM checkpoint projection; benchmark inline/precompute with the target MTP width and workload.", + help="ReplaySSM projection; parallel uses a Tensor Core GDN MTP solver (BF16, 128x128 heads).", + ) + parser.add_argument( + "--replayssm_fold_method", + choices=["raw", "derived"], + default="raw", + help="ReplaySSM fold: raw replays accepted inputs; derived reuses parallel-solver updates and changes " + "checkpoint rounding. Derived requires parallel projection; state precision and capacity are unchanged.", ) parser.add_argument( "--linear_att_ssm_data_type", diff --git a/lightllm/server/api_start.py b/lightllm/server/api_start.py index 8512baa3fc..614120cdb5 100644 --- a/lightllm/server/api_start.py +++ b/lightllm/server/api_start.py @@ -36,10 +36,17 @@ def _set_envs_and_config(args: StartArgs): def _validate_ssm_state_mode(args: StartArgs): assert args.ssm_state_mode in ("native", "compact", "replay"), "Invalid ssm_state_mode" - assert args.replayssm_projection_mode in ("inline", "precompute"), "Invalid replayssm_projection_mode" + assert args.replayssm_projection_mode in ("inline", "precompute", "parallel"), "Invalid replayssm_projection_mode" + assert args.replayssm_fold_method in ("raw", "derived"), "Invalid replayssm_fold_method" + assert ( + args.replayssm_fold_method == "raw" or args.replayssm_projection_mode == "parallel" + ), "derived fold requires parallel ReplaySSM projection" assert ( args.replayssm_projection_mode == "inline" or args.ssm_state_mode == "replay" - ), "replayssm_projection_mode=precompute requires ssm_state_mode=replay" + ), "Non-inline ReplaySSM projection requires ssm_state_mode=replay" + if args.replayssm_projection_mode == "parallel": + assert args.mtp_step > 0, "Parallel ReplaySSM requires MTP" + assert args.linear_att_ssm_data_type == "bfloat16", "Parallel ReplaySSM requires BF16 state" if args.ssm_state_mode == "native": return assert args.linear_att_ssm_data_type in ( diff --git a/lightllm/server/core/objs/start_args_type.py b/lightllm/server/core/objs/start_args_type.py index 6fc9f9bae4..82fc141746 100644 --- a/lightllm/server/core/objs/start_args_type.py +++ b/lightllm/server/core/objs/start_args_type.py @@ -245,6 +245,7 @@ class StartArgs: ssm_state_mode: str = "native" replayssm_cache_len: int = 16 replayssm_projection_mode: str = "inline" + replayssm_fold_method: str = "raw" 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/lightllm/server/router/model_infer/mode_backend/base_backend.py b/lightllm/server/router/model_infer/mode_backend/base_backend.py index 650cb8493c..98c10d1353 100644 --- a/lightllm/server/router/model_infer/mode_backend/base_backend.py +++ b/lightllm/server/router/model_infer/mode_backend/base_backend.py @@ -930,6 +930,17 @@ def _trans_req_ids_to_req_objs(self, req_ids: List[int]) -> List[InferReq]: return [g_infer_context.requests_mapping[req_id] for req_id in req_ids] def _gen_argmax_token_ids(self, model_output: ModelOutput): + if ( + model_output.logits_token_ids is None + and model_output.logits.is_cuda + and model_output.logits.dtype == torch.float32 + and model_output.logits.shape[0] > 0 + and model_output.logits.stride(1) == 1 + and model_output.logits.shape[1] >= 4096 + ): + from lightllm.common.basemodel.triton_kernel.post_process.argmax import argmax + + return argmax(model_output.logits) token_indices = torch.argmax(model_output.logits, dim=-1) if model_output.logits_token_ids is not None: # 候选排列保证分数相同时优先选择最小的全局 token ID。 diff --git a/lightllm/server/router/model_infer/mode_backend/generic_post_process.py b/lightllm/server/router/model_infer/mode_backend/generic_post_process.py index 5b29ea0510..39727ebd07 100644 --- a/lightllm/server/router/model_infer/mode_backend/generic_post_process.py +++ b/lightllm/server/router/model_infer/mode_backend/generic_post_process.py @@ -3,6 +3,7 @@ from lightllm.common.basemodel.triton_kernel.post_process.apply_penalty import apply_penalty from lightllm.common.basemodel.triton_kernel.post_process.apply_penalty_gpu_cache import apply_penalty_gpu_cache from lightllm.common.basemodel.triton_kernel.post_process.apply_invalid_token import apply_invalid_token_ids +from lightllm.common.basemodel.triton_kernel.post_process.greedy import greedy_sample from lightllm.server.router.model_infer.infer_batch import InferReq, g_infer_context from lightllm.server.router.model_infer.pin_mem_manager import g_pin_mem_manager from lightllm.utils.envs_utils import get_env_start_args @@ -23,6 +24,7 @@ def sample(logits: torch.Tensor, reqs: List[InferReq], eos_id: List[int] = [2]): skip_top_k, skip_top_p, exist_req_use_random_seed, + skip_token_penalties, ) = _get_post_sample_tensors(reqs) eos_ids = g_pin_mem_manager.gen_from_list(key="eos_ids", data=eos_id, dtype=torch.int32).cuda(non_blocking=True) @@ -66,6 +68,7 @@ def sample(logits: torch.Tensor, reqs: List[InferReq], eos_id: List[int] = [2]): b_mask_eos_reqs=b_mask_eos_reqs, eos_ids=eos_ids, sampling_params_manager=sampling_params_manager, + skip_token_penalties=skip_token_penalties, ) if has_invalid_token_ids: @@ -75,6 +78,9 @@ def sample(logits: torch.Tensor, reqs: List[InferReq], eos_id: List[int] = [2]): cu_invalid_token_num=cu_invalid_token_num, ) + if is_all_greedy and logits.is_cuda and logits.dtype == torch.float32 and logits.stride(1) == 1: + return greedy_sample(logits, b_temperatures) + logits.div_(b_temperatures.view((-1, 1))) probs = torch.softmax(logits, dim=-1) @@ -165,6 +171,7 @@ def _get_post_sample_tensors(reqs: List[InferReq]): skip_top_k = True skip_top_p = True exist_req_use_random_seed = False + skip_token_penalties = True # invalid token ids invalid_token_ids: List[int] = [] @@ -175,6 +182,11 @@ def _get_post_sample_tensors(reqs: List[InferReq]): for i, req_obj in enumerate(reqs): sample_param = req_obj.sampling_param shm_param = sample_param.shm_param + skip_token_penalties = skip_token_penalties and ( + shm_param.presence_penalty == 0.0 + and shm_param.frequency_penalty == 0.0 + and shm_param.repetition_penalty == 1.0 + ) exponential_decay_length_penalty = shm_param.exponential_decay_length_penalty.to_tuple() out_token_len = req_obj.get_cur_total_len() - req_obj.shm_req.input_len length_penalty_param.append(max(out_token_len - exponential_decay_length_penalty[0], 0)) @@ -230,4 +242,5 @@ def _get_post_sample_tensors(reqs: List[InferReq]): skip_top_k, skip_top_p, exist_req_use_random_seed, + skip_token_penalties, ) diff --git a/lightllm/server/router/model_infer/mtp_speculative/proposers/eagle_with_att.py b/lightllm/server/router/model_infer/mtp_speculative/proposers/eagle_with_att.py index 721af1c27d..0f9c4fe8f1 100644 --- a/lightllm/server/router/model_infer/mtp_speculative/proposers/eagle_with_att.py +++ b/lightllm/server/router/model_infer/mtp_speculative/proposers/eagle_with_att.py @@ -73,14 +73,27 @@ def propose_next( verify_draft_input = copy.copy(target_model_input) verify_draft_input.input_ids = target_next_token_ids verify_draft_input.mtp_draft_input_hiddens = target_model_output.mtp_collector.spec_hidden + if ( + draft_model.use_replayssm_draft_tail() + and verify_draft_input.batch_size > draft_model.args.running_max_req_size + ): + verify_draft_input.b_logits_rows = accepted_tail_rows extend_output = draft_model.forward(verify_draft_input) # 只在 req_num 行 logits 上进行 argmax，避免为未接受的 verify 行执行 # vocabulary reduction。第一列 proposal 来自每个请求的 accepted tail。 accepted_tail_output = ModelOutput( - logits=extend_output.logits.index_select(0, accepted_tail_rows), + logits=( + extend_output.logits[:req_num] + if extend_output.logits_row_stride > 1 + else extend_output.logits.index_select(0, accepted_tail_rows) + ), logits_token_ids=( - extend_output.logits_token_ids.index_select(0, accepted_tail_rows) + ( + extend_output.logits_token_ids[:req_num] + if extend_output.logits_row_stride > 1 + else extend_output.logits_token_ids.index_select(0, accepted_tail_rows) + ) if extend_output.logits_token_ids is not None else None ), diff --git a/test/benchmark/kernels/benchmark_replayssm_service.py b/test/benchmark/kernels/benchmark_replayssm_service.py new file mode 100644 index 0000000000..ea65431ad4 --- /dev/null +++ b/test/benchmark/kernels/benchmark_replayssm_service.py @@ -0,0 +1,118 @@ +"""Compare streamed requests on baseline and optimized services with identical warmup. + +Reports request timings and token IDs for an exact cross-revision comparison. +""" + +import argparse +import concurrent.futures +import json +import statistics +import time +import threading +from pathlib import Path +import requests +from transformers import AutoTokenizer + +p = argparse.ArgumentParser() +p.add_argument("--output", required=True) +p.add_argument("--batch", type=int, nargs="+", default=[1, 8, 16]) +p.add_argument("--repeats", type=int, default=3) +p.add_argument("--tokens", type=int, default=64) +p.add_argument("--warmups", type=int, default=1) +p.add_argument("--model-dir", required=True) +p.add_argument("--url", default="http://127.0.0.1:18785/generate_stream") +args = p.parse_args() +url = args.url +tok = AutoTokenizer.from_pretrained(args.model_dir, trust_remote_code=True) +prefix = tok.encode( + "Describe how a computer executes a program, from source code to machine instructions. " + "Explain the role of memory, processors, and operating systems. " * 100 +)[:512] + + +def prepare_prompt(index, rep): + suffix = tok.encode( + f"Example number {rep * 1000 + index}. Discuss this scenario in detail. " + + ("An engineer measures performance and validates numerical correctness. " * 100) + )[:512] + return prefix + suffix + + +def one(index, rep, ids, barrier, cohort_start): + payload = { + "inputs": ids, + "parameters": { + "do_sample": False, + "disable_prompt_cache": True, + "ignore_eos": True, + "max_new_tokens": args.tokens, + }, + } + barrier.wait() + start = time.perf_counter() + times = [] + tokens = [] + texts = [] + finished = False + with requests.post(url, json=payload, stream=True, timeout=300) as response: + response.raise_for_status() + # Consume each HTTP chunk; byte-sized iteration can saturate the client CPU. + for line in response.iter_lines(chunk_size=None): + if not line.startswith(b"data:"): + continue + item = json.loads(line[5:]) + assert "error" not in item, item + token = item.get("token", {}) + if token.get("id") is not None: + times.append(time.perf_counter() - cohort_start) + tokens.append(token["id"]) + texts.append(token.get("text", "")) + finished = item.get("finished", finished) + assert finished and len(tokens) == args.tokens, (finished, len(tokens)) + return dict( + index=index, + rep=rep, + prompt_len=len(ids), + token_ids=tokens, + text="".join(texts), + token_times_s=times, + ttft_ms=(times[0] - (start - cohort_start)) * 1000, + tpot_ms=(times[-1] - times[0]) * 1000 / max(1, len(times) - 1), + elapsed_s=time.perf_counter() - start, + ) + + +def run(batch, rep): + prompts = [prepare_prompt(i, rep) for i in range(batch)] + barrier = threading.Barrier(batch) + start = time.perf_counter() + with concurrent.futures.ThreadPoolExecutor(max_workers=batch) as pool: + samples = list(pool.map(lambda i: one(i, rep, prompts[i], barrier, start), range(batch))) + elapsed = time.perf_counter() - start + begin = max(r["token_times_s"][0] for r in samples) + 0.5 + end = min(r["token_times_s"][-1] for r in samples) - 0.5 + count = sum(sum(begin <= t <= end for t in r["token_times_s"]) for r in samples) + return dict( + batch=batch, + rep=rep, + elapsed_s=elapsed, + tokens_per_s=sum(len(r["token_ids"]) for r in samples) / elapsed, + steady_tokens_per_s=count / (end - begin) if end > begin else None, + steady_window_s=max(0.0, end - begin), + ttft_ms=statistics.median(r["ttft_ms"] for r in samples), + tpot_ms=statistics.median(r["tpot_ms"] for r in samples), + requests=samples, + ) + + +# Warm every measured graph bucket before timing, with the same requests on both revisions. +for batch in args.batch: + for warm in range(args.warmups): + run(batch, 99 + warm) +rows = [] +for batch in args.batch: + for rep in range(args.repeats): + row = run(batch, rep) + rows.append(row) + Path(args.output).write_text(json.dumps(rows, indent=2)) + print(json.dumps({k: v for k, v in row.items() if k != "requests"}), flush=True) diff --git a/unit_tests/common/basemodel/test_model_output.py b/unit_tests/common/basemodel/test_model_output.py index 3f6664d031..95acb17060 100644 --- a/unit_tests/common/basemodel/test_model_output.py +++ b/unit_tests/common/basemodel/test_model_output.py @@ -229,3 +229,38 @@ def forward(state): assert seen_sizes == [expected_size] * (2 if overlap else 1) torch.testing.assert_close(output.logits, expected) torch.testing.assert_close(torch.stack((model_input.b_req_idx, model_input.b_seq_len), dim=1), expected) + + +@pytest.mark.parametrize("origin", [0, 5, 10, 15]) +def test_selected_draft_logits_unpad_preserves_physical_hidden_rows(origin): + model = TpPartBaseModel.__new__(TpPartBaseModel) + output = ModelOutput( + logits=torch.arange(12).view(3, 4), + logits_token_ids=torch.arange(12).view(3, 4), + mtp_collector=ModelMtpOutputCollector(spec_hidden=torch.arange(45).view(15, 3)), + logits_row_stride=5, + ) + unpadded = model._create_unpad_decode_model_output(output, origin) + assert unpadded.logits.shape == (origin // 5, 4) + assert unpadded.logits_token_ids.shape == (origin // 5, 4) + assert unpadded.mtp_collector.spec_hidden.shape == (origin, 3) + assert output.mtp_collector.spec_hidden.shape == (15, 3) + assert unpadded.logits_row_stride == 5 + + +@pytest.mark.parametrize("disabled", [None, "dynamic", "overlap", "tp", "projection", "model", "target"]) +def test_draft_tail_scope(disabled): + from lightllm.server.core.objs.start_args_type import StartArgs + + model = TpPartBaseModel.__new__(TpPartBaseModel) + model.args = StartArgs( + ssm_state_mode="replay", replayssm_projection_mode="parallel", mtp_mode="eagle_with_att", mtp_step=4 + ) + model.is_mtp_draft_model = disabled != "target" + model.tp_world_size_ = 2 if disabled == "tp" else 1 + model.config = {"model_type": "qwen3_5_text" if disabled != "model" else "qwen3_next"} + model.args.mtp_dynamic_verify = disabled == "dynamic" + model.args.enable_decode_microbatch_overlap = disabled == "overlap" + if disabled == "projection": + model.args.replayssm_projection_mode = "inline" + assert model.use_replayssm_draft_tail() == (disabled is None) diff --git a/unit_tests/common/basemodel/triton_kernel/linear_att/test_replayssm_parallel.py b/unit_tests/common/basemodel/triton_kernel/linear_att/test_replayssm_parallel.py new file mode 100644 index 0000000000..1eafafae20 --- /dev/null +++ b/unit_tests/common/basemodel/triton_kernel/linear_att/test_replayssm_parallel.py @@ -0,0 +1,71 @@ +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") + + +@pytest.mark.parametrize("capacity", [8, 16, 64]) +def test_derived_fold_commits_only_accepted_history(capacity): + torch.manual_seed(731) + state = torch.randn(2, 5, 4, 128, 128, device="cuda", dtype=torch.bfloat16) * 0.02 + initial = state.clone() + cache = ReplaySSMCache(state, capacity, 5, num_key_heads=2, projection_mode="parallel", fold_method="derived") + cache.keys.fill_(float("nan")) + cache.deltas.fill_(float("nan")) + cache.gates.fill_(float("nan")) + cache.raw_keys.fill_(float("nan")) + cache.raw_values.fill_(float("nan")) + cache.betas.fill_(float("nan")) + # Fold request 2 from the second ring half; request 0 has an empty span. + count = capacity - 2 + cache.cursors[2] = 2 * capacity + count + cache.cursors[0] = count + cache.keys[:, 2, :, capacity : capacity + count].normal_(0, 0.1) + cache.deltas[:, 2, :, capacity : capacity + count].normal_(0, 0.1) + cache.gates[:, 2, :, capacity : capacity + count].fill_(-0.1) + expected = initial[:, 2].float() + for j in range(count): + keys = cache.keys[:, 2, :, capacity + j].repeat_interleave(2, dim=1) + delta = cache.deltas[:, 2, :, capacity + j] + decay = cache.gates[:, 2, :, capacity + j].exp() + expected = expected * decay[:, :, None, None] + keys[:, :, :, None] * delta[:, :, None, :] + reqs = torch.tensor([2, 0, 4], dtype=torch.int32, device="cuda") + cu = torch.tensor([0, 5, 5, 10], dtype=torch.int32, device="cuda") + cache.prepare_decode(reqs, cu) + assert torch.isfinite(state).all() + torch.testing.assert_close(state[:, 2], expected.bfloat16(), atol=2e-4, rtol=0.008) + torch.testing.assert_close(state[:, [0, 1, 3, 4]], initial[:, [0, 1, 3, 4]], atol=0, rtol=0) + assert cache.cursors[2] == 0 + assert cache.cursors[0] == count + assert cache.cursors[4] == 0 + + +@pytest.mark.parametrize("fold", ["raw", "derived"]) +def test_parallel_rejects_unsupported_shapes(fold): + state = torch.zeros(1, 3, 4, 32, 64, device="cuda", dtype=torch.bfloat16) + with pytest.raises(AssertionError, match="128x128"): + ReplaySSMCache(state, 16, 5, num_key_heads=2, projection_mode="parallel", fold_method=fold) + + +def test_derived_fold_requires_parallel_solver(): + state = torch.zeros(1, 3, 4, 128, 128, device="cuda", dtype=torch.bfloat16) + with pytest.raises(AssertionError): + ReplaySSMCache(state, 16, 5, num_key_heads=2, fold_method="derived") + + +@pytest.mark.parametrize("width", [2, 3, 5, 8, 16]) +def test_parallel_raw_against_sequential_reference(width): + from unit_tests.common.basemodel.triton_kernel.linear_att.test_replayssm import ( + test_replay_acceptance_flush_and_graph as check_replay, + ) + + check_replay( + width=width, + dims=(128, 128), + capacity=16, + dtype=torch.bfloat16, + projection_mode="parallel", + run_config={"BV": 128, "num_warps": 2, "num_stages": 1}, + ) diff --git a/unit_tests/common/basemodel/triton_kernel/linear_att/test_ssm_autotune.py b/unit_tests/common/basemodel/triton_kernel/linear_att/test_ssm_autotune.py index 0665b949ac..36f550a04c 100644 --- a/unit_tests/common/basemodel/triton_kernel/linear_att/test_ssm_autotune.py +++ b/unit_tests/common/basemodel/triton_kernel/linear_att/test_ssm_autotune.py @@ -350,6 +350,7 @@ def forward(target): hidden_collector=SimpleNamespace(finish_output=lambda **kw: ModelMtpOutputCollector()), prompt_logics=None, ssm_run_config=config, + logits_row_stride=1, ) output = TpPartBaseModel._create_model_output( None, PostLayerOutput(logits=outputs[0].view(count * 4, -1)), state diff --git a/unit_tests/common/basemodel/triton_kernel/test_greedy_sampling.py b/unit_tests/common/basemodel/triton_kernel/test_greedy_sampling.py new file mode 100644 index 0000000000..ee7980b3a2 --- /dev/null +++ b/unit_tests/common/basemodel/triton_kernel/test_greedy_sampling.py @@ -0,0 +1,142 @@ +import pytest +import torch + +from lightllm.common.basemodel.triton_kernel.post_process.argmax import argmax +from lightllm.common.basemodel.triton_kernel.post_process.greedy import greedy_sample + +pytestmark = pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") + + +@pytest.mark.parametrize("vocab", [13, 4101, 248320]) +def test_greedy_matches_torch_with_special_logits(vocab): + torch.manual_seed(417) + logits = torch.randn(7, vocab, device="cuda") + logits[0].fill_(-float("inf")) + logits[1, 2] = float("nan") + logits[1, 9] = float("nan") + logits[2, 4] = float("inf") + logits[3, 4] = logits[3, 9] = 100 + temperatures = torch.tensor([1, 1, 1, 0.5, 0.7, 1.3, 2.0], device="cuda") + expected = logits / temperatures[:, None] + ids = expected.argmax(-1) + probs = expected.softmax(-1).gather(1, ids[:, None]).log().flatten() + actual_ids, actual_probs = greedy_sample(logits, temperatures) + assert torch.equal(ids, actual_ids) + assert torch.equal(argmax(expected), ids) + torch.testing.assert_close(logits, expected, atol=0, rtol=0, equal_nan=True) + torch.testing.assert_close(actual_probs, probs, atol=3e-6, rtol=1e-6, equal_nan=True) + + +@pytest.mark.parametrize("vocab", [13, 4101]) +def test_identity_penalty_skip_preserves_eos_and_masks(vocab): + from types import SimpleNamespace + from lightllm.common.basemodel.triton_kernel.post_process.apply_penalty_gpu_cache import apply_penalty_gpu_cache + + torch.manual_seed(41) + manager = SimpleNamespace( + vocab_size=vocab, + req_to_presence_penalty=torch.zeros(7, device="cuda"), + req_to_frequency_penalty=torch.zeros(7, device="cuda"), + req_to_repetition_penalty=torch.ones(7, device="cuda"), + req_to_out_token_id_counter=torch.randint(0, 20, (7, vocab), device="cuda", dtype=torch.int32), + req_to_exponential_decay_length_penalty=torch.tensor([1.0, 1.01, 1.2, 1.0, 1.5, 0.9, 1.0], device="cuda"), + ) + logits = torch.randn(5, vocab, device="cuda") + original = logits.clone() + params = dict( + b_req_idx=torch.tensor([5, 3, 0, 4, 1], device="cuda", dtype=torch.int32), + b_length_penalty_param=torch.tensor([8, 0, 3, 4, 2], device="cuda", dtype=torch.int32), + b_mask_eos_reqs=torch.tensor([False, True, False, True, False], device="cuda"), + eos_ids=torch.tensor([2, 9], device="cuda", dtype=torch.int32), + sampling_params_manager=manager, + ) + apply_penalty_gpu_cache(original, **params) + apply_penalty_gpu_cache(logits, **params, skip_token_penalties=True) + torch.testing.assert_close(logits, original, rtol=0, atol=0) + assert (logits[[1, 3]][:, [2, 9]] == -10000000).all() + + +@pytest.mark.parametrize("greedy", [True, False]) +@pytest.mark.parametrize("identity", [True, False]) +def test_full_sampler_with_penalties_invalid_ids_and_random_sampling(monkeypatch, greedy, identity): + from types import SimpleNamespace + import lightllm.server.router.model_infer.mode_backend.generic_post_process as post + + vocab, batch = 4101, 4 + manager = SimpleNamespace( + penalty_counter_mode="gpu_counter", + vocab_size=vocab, + req_to_presence_penalty=torch.full((batch,), 0.0 if identity else 0.5, device="cuda"), + req_to_frequency_penalty=torch.full((batch,), 0.0 if identity else 0.1, device="cuda"), + req_to_repetition_penalty=torch.full((batch,), 1.0 if identity else 1.2, device="cuda"), + req_to_out_token_id_counter=torch.randint(0, 10, (batch, vocab), device="cuda", dtype=torch.int32), + req_to_exponential_decay_length_penalty=torch.full((batch,), 1.01, device="cuda"), + ) + reqs = [] + for index in range(batch): + param = SimpleNamespace( + presence_penalty=0.0 if identity else 0.5, + frequency_penalty=0.0 if identity else 0.1, + repetition_penalty=1.0 if identity else 1.2, + temperature=0.7 + 0.2 * index, + top_k=1 if greedy else vocab, + top_p=1.0, + min_new_tokens=12 if index % 2 else 0, + exponential_decay_length_penalty=SimpleNamespace(to_tuple=lambda: (0, 1.01)), + ) + reqs.append( + SimpleNamespace( + sampling_param=SimpleNamespace(shm_param=param, invalid_token_ids=[0, 5, 13]), + shm_req=SimpleNamespace(input_len=3), + get_cur_total_len=lambda: 10, + vocab_size=vocab, + generator=None, + req_idx=index, + ) + ) + monkeypatch.setattr( + post, "g_infer_context", SimpleNamespace(req_manager=SimpleNamespace(req_sampling_params_manager=manager)) + ) + monkeypatch.setattr( + post, + "g_pin_mem_manager", + SimpleNamespace(gen_from_list=lambda key, data, dtype: torch.tensor(data, dtype=dtype)), + ) + torch.manual_seed(62) + logits = torch.randn(batch, vocab, device="cuda") + torch.manual_seed(31) + actual = post.sample(logits.clone(), reqs, eos_id=[2, 9]) + + original_penalty = post.apply_penalty_gpu_cache + + def unskipped_penalty(**kwargs): + kwargs["skip_token_penalties"] = False + return original_penalty(**kwargs) + + def reference_greedy(logits, temperatures): + logits.div_(temperatures[:, None]) + probs = logits.softmax(-1) + ids = logits.argmax(-1) + return ids, probs.gather(1, ids[:, None]).log().flatten() + + monkeypatch.setattr(post, "apply_penalty_gpu_cache", unskipped_penalty) + monkeypatch.setattr(post, "greedy_sample", reference_greedy) + torch.manual_seed(31) + expected = post.sample(logits.clone(), reqs, eos_id=[2, 9]) + torch.testing.assert_close(actual[0], expected[0], rtol=0, atol=0) + torch.testing.assert_close(actual[1], expected[1], rtol=1e-6, atol=3e-6) + + +def test_greedy_fully_masked_partition_keeps_finite_logprob(): + logits = torch.randn(3, 8201, device="cuda") + logits[0, :4096] = -float("inf") + logits[1, 4096:8192] = -float("inf") + logits[2, 8192:] = -float("inf") + temperatures = torch.tensor([0.7, 1.0, 1.3], device="cuda") + scaled = logits / temperatures[:, None] + expected_ids = scaled.argmax(-1) + expected_logp = scaled.softmax(-1).gather(1, expected_ids[:, None]).log().flatten() + ids, logp = greedy_sample(logits, temperatures) + torch.testing.assert_close(ids, expected_ids, rtol=0, atol=0) + torch.testing.assert_close(logp, expected_logp, rtol=1e-6, atol=3e-6) + assert torch.isfinite(logp).all() diff --git a/unit_tests/server/router/model_infer/mtp_speculative/test_eagle_with_att.py b/unit_tests/server/router/model_infer/mtp_speculative/test_eagle_with_att.py index 365acb30b2..894833ab03 100644 --- a/unit_tests/server/router/model_infer/mtp_speculative/test_eagle_with_att.py +++ b/unit_tests/server/router/model_infer/mtp_speculative/test_eagle_with_att.py @@ -23,6 +23,7 @@ def forward(model_input): draft_model = SimpleNamespace( forward=forward, + use_replayssm_draft_tail=lambda: False, map_draft_vocab_to_main_vocab=lambda token_ids: token_ids + 100, ) proposer = Eagle3Proposer( @@ -102,7 +103,9 @@ def forward(model_input): return ModelOutput(logits=torch.empty((0,), device=device)) proposer = EagleWithAttProposer( - backend=SimpleNamespace(draft_models=[SimpleNamespace(forward=forward)]), + backend=SimpleNamespace( + draft_models=[SimpleNamespace(forward=forward, use_replayssm_draft_tail=lambda: False)] + ), enable_dynmaic_mtp=False, ) proposer.fill_draft_model_kv_state( @@ -158,7 +161,7 @@ def forward(model_input): ) backend = SimpleNamespace( - draft_models=[SimpleNamespace(forward=forward)], + draft_models=[SimpleNamespace(forward=forward, use_replayssm_draft_tail=lambda: False)], _gen_argmax_token_ids=lambda output: output.logits[:, 0].long(), _gen_argmax_token_ids_and_prob=lambda output: ( output.logits[:, 0].long(), @@ -211,3 +214,43 @@ def forward(model_input): assert draft_calls[1]["total_token_num"] == 46 assert target_input.input_ids is original_input_ids assert target_input.mtp_draft_input_hiddens is None + + +def test_selected_tail_logits_keep_complete_draft_kv_input(): + hidden = torch.arange(20).view(10, 2).float() + target_input = SimpleNamespace( + is_prefill=False, + batch_size=10, + input_ids=None, + b_position_delta=torch.zeros(10, dtype=torch.int32), + mtp_draft_input_hiddens=None, + b_logits_rows=None, + ) + + def forward(model_input): + assert model_input.batch_size == 10 + assert model_input.mtp_draft_input_hiddens is hidden + assert model_input.input_ids.numel() == 10 + torch.testing.assert_close(model_input.b_logits_rows, torch.tensor([1, 8])) + return ModelOutput(logits=torch.tensor([[7.0], [9.0]]), logits_row_stride=5) + + draft_model = SimpleNamespace( + forward=forward, + use_replayssm_draft_tail=lambda: True, + args=SimpleNamespace(running_max_req_size=2), + ) + proposer = EagleWithAttProposer( + backend=SimpleNamespace(draft_models=[draft_model], _gen_argmax_token_ids=lambda out: out.logits[:, 0].long()), + enable_dynmaic_mtp=False, + ) + output = proposer.propose_next( + target_input, + ModelOutput(logits=torch.empty(10, 1), mtp_collector=ModelMtpOutputCollector(spec_hidden=hidden)), + torch.arange(10), + torch.tensor([0, 5]), + 1, + torch.tensor([2, 4]), + ) + torch.testing.assert_close(output.token_ids, torch.tensor([[7], [9]])) + assert target_input.b_logits_rows is None + assert target_input.input_ids is None diff --git a/unit_tests/server/test_ssm_state_mode.py b/unit_tests/server/test_ssm_state_mode.py index b716850207..6016cd2a31 100644 --- a/unit_tests/server/test_ssm_state_mode.py +++ b/unit_tests/server/test_ssm_state_mode.py @@ -83,3 +83,30 @@ def test_replay_projection_mode_requires_replay(monkeypatch): args.ssm_state_mode = "compact" with pytest.raises(AssertionError, match="requires ssm_state_mode=replay"): _validate_ssm_state_mode(args) + + +@pytest.mark.parametrize("fold", ["raw", "derived"]) +def test_parallel_replay_configuration(monkeypatch, fold): + from lightllm.server.core.objs.start_args_type import StartArgs + + monkeypatch.setattr("lightllm.utils.config_utils.get_model_type", lambda _: "qwen3_5") + args = StartArgs( + ssm_state_mode="replay", + replayssm_projection_mode="parallel", + replayssm_fold_method=fold, + mtp_step=4, + mtp_mode="eagle_with_att", + linear_att_ssm_data_type="bfloat16", + ) + _validate_ssm_state_mode(args) + args.mtp_step = 0 + with pytest.raises(AssertionError, match="requires MTP"): + _validate_ssm_state_mode(args) + + +def test_derived_replay_requires_parallel_configuration(): + from lightllm.server.core.objs.start_args_type import StartArgs + + args = StartArgs(ssm_state_mode="replay", replayssm_fold_method="derived") + with pytest.raises(AssertionError, match="requires parallel"): + _validate_ssm_state_mode(args) From b58019818bb2e88b02955b33c39e2960b24144c1 Mon Sep 17 00:00:00 2001 From: sufubao Date: Wed, 14 Oct 2026 14:51:47 +0800 Subject: [PATCH 11/40] style: format ReplaySSM kernels --- .../common/basemodel/triton_kernel/linear_att/replayssm.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/lightllm/common/basemodel/triton_kernel/linear_att/replayssm.py b/lightllm/common/basemodel/triton_kernel/linear_att/replayssm.py index 17c3a1cfc5..3db7cf2f98 100644 --- a/lightllm/common/basemodel/triton_kernel/linear_att/replayssm.py +++ b/lightllm/common/basemodel/triton_kernel/linear_att/replayssm.py @@ -284,7 +284,7 @@ def _replay( qk /= tl.sqrt(tl.sum(qk * qk, axis=1) + 1.0e-6)[:, None] else: qk *= tl.rsqrt(tl.sum(qk * qk, axis=1) + 1.0e-6)[:, None] - qk *= tl.where(tt < BW, K**-0.5, 1.0)[:, None] + qk *= tl.where(tt < BW, K ** -0.5, 1.0)[:, None] if KDA: history = ll < n gates = tl.load( @@ -467,7 +467,7 @@ def _replay( 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) + 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) if KDA: x = tl.load(A + t * SA + hv * K + kk, kk < K, 0).to(tl.float32) + bias From 62258cdaa0b7eb5d56774d8d59f22fc642a8fcfc Mon Sep 17 00:00:00 2001 From: sufubao Date: Wed, 14 Oct 2026 14:54:02 +0800 Subject: [PATCH 12/40] chore: exclude ReplaySSM docs from PR --- docs/replayssm/README.md | 133 ------------------------ docs/replayssm/h200_optimizations.md | 125 ---------------------- docs/replayssm/h200_retest_summary.json | 86 --------------- 3 files changed, 344 deletions(-) delete mode 100644 docs/replayssm/README.md delete mode 100644 docs/replayssm/h200_optimizations.md delete mode 100644 docs/replayssm/h200_retest_summary.json diff --git a/docs/replayssm/README.md b/docs/replayssm/README.md deleted file mode 100644 index 27cb5e5678..0000000000 --- a/docs/replayssm/README.md +++ /dev/null @@ -1,133 +0,0 @@ -# ReplaySSM for GDN and KDA - -ReplaySSM reduces speculative SSM storage by keeping one checkpoint per request -and replaying accepted updates. It is opt-in; `--ssm_state_mode native` remains -the default. - -## Configuration - -```bash -LIGHTLLM_TRITON_AUTOTUNE_LEVEL=1 python -m lightllm.server.api_server \ - --model_dir /path/to/qwen3.5 \ - --ssm_state_mode replay \ - --linear_att_ssm_data_type bfloat16 \ - --replayssm_cache_len 8 \ - --replayssm_projection_mode precompute -``` - -The existing MTP flags can be added to this command. History capacity must cover -the maximum verify width (`mtp_step + 1`), including dynamic MTP. Supported -capacities are 4, 8, 16, 32 and 64. - -| State mode | Storage and update policy | -| --- | --- | -| `native` | Existing per-token speculative state snapshots. | -| `compact` | One state plus this round's raw inputs; merge the accepted prefix every round. Requires MTP. | -| `replay` | One state plus bounded accepted history; fold before the next verify window would overflow. Supports ordinary decode and MTP. | - -Non-native modes support Qwen3-Next/Qwen3.5 GDN and GLM-5.3-FLASH KDA with FP32 or BF16 -state. Other models retain their native path. `--replayssm_projection_mode` -selects `inline` (default), `precompute`, or `parallel`; the latter two require replay. -`parallel` uses a Tensor Core triangular solve for BF16 GDN MTP with 128x128 -heads. It currently uses the measured BV128 / 2-warp configuration. -Choose the projection mode and capacity using the target workload. -Without MTP, replay adds history storage to the single native checkpoint. - -### Qwen3.5 H200 experimental optimization - -For the measured Qwen3.5-27B / TP1 / BF16 / MTP4 configuration: - -```bash ---ssm_state_mode replay --replayssm_cache_len 16 \ ---replayssm_projection_mode parallel --replayssm_fold_method derived -``` - -`--replayssm_fold_method raw` remains the default. `derived` requires `parallel` -and folds cached FP32 keys/deltas with FP32 matrix accumulation. It changes -checkpoint rounding relative to raw-input replay and is not bitwise equivalent. -Neither mode changes the configured request, KV, or history capacity. - -For Qwen3.5 dense TP1 with fixed EAGLE-with-attention MTP, `parallel` also -captures first-draft graphs above the logical request capacity and projects only -accepted-tail hidden rows. All draft KV writes and collected hidden rows remain -present. Dynamic MTP, decode overlap, and other models retain the existing draft -layout. Expanded graphs consume additional graph memory; capacity settings are -unchanged, but this is not a claim of identical total GPU memory usage. - -See [H200 measurements](h200_optimizations.md) for paired -throughput, full GSM8K results, and the numerical tradeoff. Performance depends on -batch size and input/output length; the main-scene gain is not a universal bound. - -## State lifecycle and precision - -Verify overwrites the uncommitted history suffix. Acceptance advances only the -accepted prefix, so rejected proposals never enter a checkpoint. Fold alternates -history buffers to avoid overwriting records still read by another CTA. Padding -and the HOLD request do not advance history. - -Prefill folds accepted updates before using the existing chunk kernel. CPU -prefix checkpoints and PD pages contain canonical SSM state and the accepted -conv window, without replay scratch. Exporting a checkpoint does not mutate -the active state. Restoring a checkpoint or reusing a request slot invalidates -history and resets the MTP offset. - -Replay accumulates in FP32 and rounds when folding or materializing a checkpoint. -BF16 native and compact modes round after each token. Replay therefore does not -promise bitwise equivalence or identical generated text to native BF16. The -`precompute` path also uses activation-precision projections. Correctness tests -use an independent recurrence with each mode's rounding boundary. - -The SSM autotuner measures prepare, forward, acceptance and fold on disposable -state; serving state is never a tuning input. Each batch bucket selects its own -layout before CUDA Graph capture. Model outputs carry that layout through -unpadding, and compact acceptance uses the corresponding verify layout for each -microbatch. Replay checkpoint reconstruction keeps a batch-independent layout. -Cached configurations distinguish token count, sequence capacity, precision, -stride, heads, verify width, history capacity, projection mode and recurrent rule. - -KDA uses a per-key-dimension log decay, while GDN uses one scalar per value -head. KDA history stores the full gate vector; checkpoint projection applies -the cumulative vector decay to Q/K before projecting. The low-rank correction -uses forward decay between accepted tokens, without inverse exponential -scaling. The model's `gate_lower_bound` also separates autotune configurations. -Accepted-state exports retain GLM's indexer tail for PD; CPU prefix restore -clears the tail at aligned pool boundaries. - -## Memory and performance - -See [H200 / Qwen3.5-27B optimization measurements](h200_optimizations.md) -for workload coverage, numerical tradeoffs and benchmark commands. - -For a TP4 model with 48 GDN layers, 12 local value heads and 128x128 BF16 state, -one full-model checkpoint occupies 18 MiB per GPU. Native MTP3 uses 72 MiB per -request. Replay L8/precompute uses about 24.07 MiB including history and cursors. -Conv state, full-attention KV, request tables and sampling buffers are additional -allocations. - -For GLM-5.3-FLASH TP8 (34 KDA layers, 8 local heads), one FP32 checkpoint -is 17 MiB per GPU. Native MTP2 uses 51 MiB per request; replay L8/inline -uses about 25.52 MiB including history. With BF16 state the corresponding -figures are 25.5 and 17.02 MiB. - -More KV space improves throughput only when the workload and request capacity -can use it. Increasing request capacity can also increase per-token latency. -Keep algorithm comparisons at equal capacity separate from capacity comparisons. - -## Tests - -```bash -exp -m "ReplaySSM GDN/KDA lifecycle regression" /usr/bin/env PYTHONPATH=. python -m pytest -q \ - unit_tests/server/test_ssm_state_mode.py \ - unit_tests/common/basemodel/attention/linear/test_gdn.py \ - unit_tests/common/basemodel/triton_kernel/linear_att/test_replayssm*.py \ - unit_tests/common/basemodel/triton_kernel/linear_att/test_ssm_autotune.py \ - unit_tests/common/basemodel/triton_kernel/linear_att/test_mtp_state_params.py \ - unit_tests/common/basemodel/triton_kernel/linear_att/test_acceptance_fusion.py \ - unit_tests/common/test_paged_kv_transfer.py \ - unit_tests/models/glm5_next/test_{mtp,cache,pd_cache}.py -``` - -CUDA tests cover partial acceptance, history folds, BF16 rounding, grouped keys, -dynamic verify lengths, CUDA Graph replay, HOLD padding, checkpoint/PD transfer, -slot reuse and fused acceptance. CPU tests cover startup validation and tuner -isolation. CUDA tests skip when a GPU is unavailable. diff --git a/docs/replayssm/h200_optimizations.md b/docs/replayssm/h200_optimizations.md deleted file mode 100644 index ac034ee911..0000000000 --- a/docs/replayssm/h200_optimizations.md +++ /dev/null @@ -1,125 +0,0 @@ -# H200 / Qwen3.5-27B：负载覆盖与正常入口验证 - -2026-10-13。固定原 PR `3af33c2f47f1738e17eb70a3c1e8f42e25ed65ed`、TP1、MTP4、BF16 SSM、L16、264 请求容量和 65536 KV token。所有输入加输出不超过 4096 token。没有通过增加资源容量或更改 MTP 配置获得收益。 - -本次包含 GPU2 的 12 组负载、72 轮交替配对（144 个 cohort），以及 GPU5 正常服务入口的独立整合验证。正常入口由 `python -m lightllm.server.api_server` 启动；候选额外启用 `--replayssm_projection_mode parallel --replayssm_fold_method derived`。 - -## 负载矩阵：同卡交替配对 - -每组、每臂预热两轮，再进行六轮正反顺序交替。greedy、关闭 thinking、禁用 prompt cache 查询、固定输出长度并 ignore_eos。相同 rep 使用相同提示。短输入使用四类技术解释模板及变化的示例编号；长输入使用固定 filler，不能代表所有生产内容。吞吐包含 prefill、请求提交和尾部完成，按总输出 token / 总 wall time 计算。95%区间为10000次配对 bootstrap，seed=417；没有删掉较慢轮次。 - -| 输入 / 输出 token | 并发 | 原 PR token/s | 候选 token/s | 增幅 | 95%区间 | -| --- | ---: | ---: | ---: | ---: | --- | -| 短输入 / 384 | 1 | 140.8 | 170.4 | 21.01% | [18.85%, 23.11%] | -| 短输入 / 384 | 8 | 961.6 | 1059.4 | 10.17% | [5.97%, 14.51%] | -| 短输入 / 384 | 32 | 2551.0 | 2956.0 | 15.88% | [13.77%, 18.30%] | -| 短输入 / 384 | 52 | 2680.6 | 3144.7 | 17.31% | [15.00%, 19.74%] | -| 短输入 / 384 | 53 | 2437.2 | 3177.4 | 30.37% | [27.69%, 33.46%] | -| 短输入 / 384 | 64 | 3083.4 | 3669.0 | 18.99% | [17.51%, 20.16%] | -| 短输入 / 384 | 128 | 3655.0 | 4467.6 | 22.23% | [20.58%, 23.66%] | -| 短输入 / 128 | 32 | 2218.9 | 2465.0 | 11.09% | [5.53%, 14.76%] | -| 短输入 / 128 | 64 | 2585.1 | 2991.6 | 15.72% | [8.73%, 22.32%] | -| 短输入 / 128 | 128 | 3077.0 | 3639.8 | 18.29% | [13.34%, 22.48%] | -| 1024 / 384 | 32 | 1558.6 | 1721.5 | 10.46% | [8.53%, 12.00%] | -| 3072 / 384 | 16 | 705.8 | 741.3 | 5.04% | [3.40%, 7.16%] | - -B128 / 输出384的主场景再次超过20%；其他负载不能套用这一数字。B52→B53跨过首轮draft的原图覆盖边界，组合收益从17.3%变为30.4%，与图扩展/accepted-tail投影机制一致，但该比较不是只切换一个组件的消融。输出缩短或长输入占比增加，端到端收益降低。 - -## 延迟 - -| 输入 / 输出 | 并发 | TTFT p50 原PR→候选 ms | TPOT p50 原PR→候选 ms | -| --- | ---: | ---: | ---: | -| 短 / 384 | 1 | 142.8 → 141.4 | 6.74 → 5.52 | -| 短 / 384 | 8 | 219.8 → 163.9 | 7.24 → 6.71 | -| 短 / 384 | 32 | 370.6 → 383.6 | 10.47 → 8.88 | -| 短 / 384 | 52 | 590.8 → 570.8 | 16.85 → 14.41 | -| 短 / 384 | 53 | 600.7 → 534.1 | 19.20 → 14.51 | -| 短 / 384 | 64 | 633.9 → 602.2 | 18.13 → 15.01 | -| 短 / 384 | 128 | 974.1 → 944.8 | 31.04 → 25.06 | -| 短 / 128 | 32 | 359.9 → 367.5 | 10.64 → 9.12 | -| 短 / 128 | 64 | 594.8 → 600.4 | 18.77 → 15.53 | -| 短 / 128 | 128 | 946.6 → 922.3 | 32.73 → 26.42 | -| 1024 / 384 | 32 | 2185.1 → 2218.5 | 13.15 → 11.31 | -| 3072 / 384 | 16 | 4070.9 → 4021.2 | 11.19 → 10.32 | - -延迟是客户端观测值，按同一负载的请求合并取分位数；TPOT为每请求首尾token时间差除以后续token数。TTFT没有一致改善。逐轮数据保留在本机 `/tmp/replayssm-h200/coverage/`。 - -## 正常服务入口 - -正常入口两臂均预热两轮并正式运行八轮。GPU5原 PR **3667.2 token/s** → 工作树候选 **4446.3 token/s**，增幅 **21.25%**，独立bootstrap95%区间 **[18.96%, 23.66%]**。候选完成后停止服务，再在同一GPU5加载原PR；这是顺序独立运行，不是交替配对。区间反映轮次波动，不覆盖两次运行之间的系统漂移。GPU2的交替配对用于补充因果证据，不能混为同一个统计设计。 -正常入口TTFT p50 原PR 913.2 → 候选 1031.2 ms；TPOT 31.21 → 24.88 ms。正常入口正式性能阶段没有其他模型客户端并发压测。 - -候选已经通过B1/B52/B53/B128的正常请求 smoke。HTTP 随机采样、top-k/top-p、非默认presence/frequency/repetition惩罚、EOS最短长度与无效token参数请求全部完成，logprob有限。Prefix checkpoint 两次实际命中1536 token，输出64 token与禁用查询的请求逐项一致。原始脚本中的 `cache_cold` 标签并非真正冷缓存：禁用查询的首次请求仍填充了缓存，以实际命中长度为准。 - -## 质量与数值权衡 - -同样的全量GSM8K1319题，greedy、关闭thinking、并发64、最大输出1536：正常入口候选 **1278/1319**，历史原 PR **1279/1319**；截断8/6。原PR独对3题、候选独对2题，879条完整token序列一致。准确率差 −0.0758个百分点，配对bootstrap95%区间 [−0.3791,+0.2274] 个百分点。这是与保存的历史质量运行比较，不是本次同卡交替质量实验，不能证明严格无损。此前实验入口 derived 为1281/1319；正常入口和实验入口1162条完整序列一致。 - -对4道有变化的题另做小并发诊断，1536上限1/4正确、截断2；3072上限2/4正确、截断1。原PR在同样小并发复测中，1536上限2/4正确、截断1，3072上限3/4正确、无截断。该诊断改变了并发，且是事后选择的困难题，只说明部分截断可以缓解；没有替换全量测试分数，也不能将差异全部归因于输出上限。 - -并行求解采用activation精度的Tensor Core投影；derived fold复用现有FP32 keys/deltas、以FP32矩阵累加合并历史。历史容量和精度未降低，但运算顺序与checkpoint舍入变化，不能保证与原PR逐位一致。默认仍为raw fold，derived显式启用；固定共享输入的长期测试覆盖 L16/8192 步及 L64/4096 步；这不等同于模型反馈下的质量保证。 - -## 保留的实现与清理范围 - -保留 parallel 求解、可选 derived fold、首轮 draft 图扩展与 accepted-tail 投影、LM head 布局、strided gate、融合 greedy、draft argmax 和默认 token 惩罚跳过。单独的 parallel 求解已测得约 11.4% 端到端提升；其余组件的组合进一步改善吞吐，不能把组合收益算作每个组件的独立收益。 - -删除批量 checkpoint、普通 decode prepare/accept 融合，以及早期特定形状的 H200 调参文件：前两项没有可靠的端到端收益，后者不用于已验证的 parallel 固定布局。恢复普通路径的 autotune 版本，保留 parallel/fold 配置隔离。旧实验原型、重复源码快照、日志、图表和过时报告移出仓库。 - -以上性能数字记录的是清理前的完整候选，不能冒充清理后的重新测量。删除的普通 decode 融合不用于 MTP4 verify；批量 checkpoint 对已测场景收益可忽略。清理后的验证针对状态、checkpoint、CUDA Graph、MTP 和采样回归。 - -## 清理后的验证 - -本地 H200 回归覆盖 676 个用例：首轮 652 通过，24 个 autotune 图隔离参数组合因旧测试桩缺少 `logits_row_stride` 而失败；补齐该字段后，autotune 与最终版采样的 74 项复测全部通过，没有未解决的失败。覆盖普通及 MTP 状态生命周期、checkpoint、GDN/KDA、CUDA Graph、parallel verify width 2/3/5/8/16、derived fold、模型输出裁剪和采样边界。计数包含复测重叠，不相加。 - -24 个修改/新增 Python 文件通过语法检查和仓库 flake8 规则，`git diff --check` 通过。10 月 13 日清理时未重跑端到端模型吞吐或 GSM8K；上述性能和质量数字均来自清理前实验。清理后的吞吐复测见下一节。日志位于本机 `/tmp/replayssm-h200/cleanup-regression.log` 与 `cleanup-recheck.log`。 - -## 2026-10-14：清理后吞吐复测 - -同一 H200 GPU5、Qwen3.5-27B、TP1、MTP4、BF16 SSM、L16、264 请求容量、65536 KV token。并发 128，每请求固定生成 384 token，沿用上一轮短输入技术解释模板，greedy、关闭 thinking、禁用 prompt cache 查询、ignore_eos。两边都通过正常 API 入口运行，各预热 2 轮、正式测量 12 轮。原 PR 仍为 `3af33c2f47f1738e17eb70a3c1e8f42e25ed65ed`；候选为清理后的代码快照，显式启用 `parallel + derived`。 - -| 指标 | 原 PR | 清理后候选 | -| --- | ---: | ---: | -| 端到端吞吐，token/s | 3663.12 | 4505.88 | -| TTFT p50，ms | 928.97 | 901.96 | -| TPOT p50，ms | 31.10 | 24.91 | -| 正式请求数 | 1536 | 1536 | -| 输出 token 数 | 589824 | 589824 | -| prompt cache 命中请求 | 0 | 0 | -| 性能窗最高温度，°C | 63 | 62 | -| 性能窗显存峰值，MiB | 104887 | 104139 | - -**吞吐提升 23.01%，独立 bootstrap 95% 区间 [21.26%, 24.52%]**（10000 次重采样，seed=417）。按全部输出 token / 全部 cohort wall time 汇总，不删除慢轮次。原 PR 的单轮范围为 3572–3741 token/s，候选为 4285–4597 token/s。全部 1179648 个输出 token 的 logprob 有限，两边都没有 HW/SW 热降频。每 verify step 输出 token 数为 3.1833 / 3.1773，候选未通过提高接受率获得这一收益。 - -这是先基线、后候选的同卡顺序独立运行，**不是交替配对**。区间反映轮次波动，不覆盖两个运行阶段之间的系统漂移；本次 23.01% 与上轮 21.25% 的差异不能归因于代码清理本身。本次仅复测这一主场景，没有重新运行其他负载或 GSM8K。结果表明移除低收益路径后，该场景的吞吐提升仍保留。 - -两份源码均已冻结并保存 Python 文件哈希；候选快照与测试结束时的工作区一致。启动参数除代码路径、优化开关和端口外一致。候选首次因旧 NCCL 端口占用而未启动成功，改用空闲端口后完成全部测量；失败启动没有产生性能样本。测试服务均已停止。 - -原始逐轮响应、客户端、启动命令、源码哈希、硬件记录和统计脚本保存在本机 `/tmp/replayssm-h200/cleanup_retest_20261014/`；[紧凑统计结果](h200_retest_summary.json)保留逐轮吞吐与汇总值。 - -## 复现与原始记录 - -正常入口使用相同模型、MTP 配置和容量，候选额外加 `--replayssm_projection_mode parallel --replayssm_fold_method derived`： - -```bash -LIGHTLLM_TRITON_AUTOTUNE_LEVEL=3 LOADWORKER=4 python -m lightllm.server.api_server \ - --model_dir /models/Qwen3.5-27B --tp 1 --host 127.0.0.1 --port 19131 --nccl_port 19132 \ - --disable_vision --disable_audio --ssm_state_mode replay --replayssm_cache_len 16 \ - --linear_att_ssm_data_type bfloat16 --llm_prefill_att_backend triton triton \ - --llm_decode_att_backend fa3 triton --graph_max_batch_size 1280 \ - --running_max_req_size 264 --max_req_total_len 4096 --max_total_token_num 65536 \ - --batch_max_tokens 4096 --chunked_prefill_size 512 --linear_att_hash_page_size 512 \ - --linear_att_page_block_num 2 --linear_att_cache_size 16 --mtp_mode eagle_with_att --mtp_step 4 \ - --replayssm_projection_mode parallel --replayssm_fold_method derived -``` - -通用 HTTP benchmark 保留在 `test/benchmark/kernels/benchmark_replayssm_service.py`，用于固定 1024-token 输入的跨版本对照；它的提示与上述短输入实验不同，不能直接复现同一吞吐数字。两臂必须使用相同参数和预热次数。 - -```bash -python test/benchmark/kernels/benchmark_replayssm_service.py \ - --model-dir /models/Qwen3.5-27B --url http://127.0.0.1:19131/generate_stream \ - --batch 32 --tokens 384 --warmups 2 --repeats 8 --output /tmp/replayssm-service.json -``` - -原始实验客户端、逐轮结果和启动命令在本机 `/tmp/replayssm-h200/{coverage,integrated,normal_baseline}/`。清理前代码及文档完整归档于 `/tmp/replayssm-h200/cleanup-before.tar.gz`。这些是本地实验记录，不是仓库依赖。 - -GPU2 矩阵性能窗最高 60°C，正常入口候选/原 PR 的 GPU5 最高 62/63°C；均无硬件或软件热降频。GPU2 部分测试与其他 GPU 的模型加载/质量测试共享主机。正常入口性能阶段没有其他模型客户端并发压测。正常入口显存峰值候选 104921 MiB、原 PR 104785 MiB；扩展图使候选多占约 136 MiB，容量配置相同不代表总显存完全相等。 diff --git a/docs/replayssm/h200_retest_summary.json b/docs/replayssm/h200_retest_summary.json deleted file mode 100644 index f205d2024e..0000000000 --- a/docs/replayssm/h200_retest_summary.json +++ /dev/null @@ -1,86 +0,0 @@ -{ - "design": "GPU5, normal API entry, baseline then candidate sequential independent runs; 2 warmups and 12 measured cohorts per arm", - "baseline_commit": "3af33c2f47f1738e17eb70a3c1e8f42e25ed65ed", - "batch": 128, - "output_tokens": 384, - "arms": { - "baseline": { - "repeats": 12, - "requests": 1536, - "output_tokens": 589824, - "tokens_per_s": 3663.117912349209, - "cohort_tokens_per_s": [ - 3572.39450013313, - 3679.1945250855033, - 3668.5673855693176, - 3664.162821283357, - 3655.0134290105893, - 3658.030321178895, - 3741.2360562287436, - 3733.451128039074, - 3630.357737495093, - 3637.8124518110835, - 3689.765095342039, - 3633.695617211311 - ], - "finite_logprobs": true, - "ttft_p50_ms": 928.9662055671215, - "tpot_p50_ms": 31.098604479478794, - "hardware": { - "samples": 139, - "max_temperature_C": 63.0, - "max_memory_MiB": 104887.0, - "min_sm_clock_MHz": 1665.0, - "max_sm_clock_MHz": 1980.0, - "thermal_slowdown_samples": 0 - }, - "acceptance": { - "logged_requests": 1536, - "prompt_cache_hits": 0, - "output_per_verify": 3.1832994219777966 - } - }, - "candidate": { - "repeats": 12, - "requests": 1536, - "output_tokens": 589824, - "tokens_per_s": 4505.876589509283, - "cohort_tokens_per_s": [ - 4401.733921373901, - 4558.013559046327, - 4454.915305507067, - 4543.451463667397, - 4460.422265583761, - 4284.845559058763, - 4583.72309215793, - 4597.109123661394, - 4551.504440860731, - 4579.715485515799, - 4497.592803040789, - 4579.283607210278 - ], - "finite_logprobs": true, - "ttft_p50_ms": 901.9625158980489, - "tpot_p50_ms": 24.913206290424494, - "hardware": { - "samples": 112, - "max_temperature_C": 62.0, - "max_memory_MiB": 104139.0, - "min_sm_clock_MHz": 1500.0, - "max_sm_clock_MHz": 1980.0, - "thermal_slowdown_samples": 0 - }, - "acceptance": { - "logged_requests": 1536, - "prompt_cache_hits": 0, - "output_per_verify": 3.17729762924417 - } - } - }, - "gain_pct": 23.006594309152373, - "independent_bootstrap_ci95_pct": [ - 21.26323367507323, - 24.51623928708578 - ], - "caveat": "The interval measures round-to-round variation, not drift between sequential service runs. This benchmark does not re-evaluate model quality." -} From 7c25a7ee60e1ba977a68759a39cfe4de99018e69 Mon Sep 17 00:00:00 2001 From: sufubao Date: Wed, 14 Oct 2026 18:55:17 +0800 Subject: [PATCH 13/40] fix: pass required arguments to SGL FA3 vision attention --- .../common/basemodel/attention_vit/fa3/fp.py | 2 ++ .../vit/triton_kernel/flashattention_nopad.py | 2 ++ .../vit/test_flash_attention_forward.py | 30 +++++++++++++++++-- 3 files changed, 31 insertions(+), 3 deletions(-) diff --git a/lightllm/common/basemodel/attention_vit/fa3/fp.py b/lightllm/common/basemodel/attention_vit/fa3/fp.py index f804116f1f..a8e56417c5 100644 --- a/lightllm/common/basemodel/attention_vit/fa3/fp.py +++ b/lightllm/common/basemodel/attention_vit/fa3/fp.py @@ -53,6 +53,8 @@ def _vit_att_fwd( pack_gqa=None, sm_margin=0, sinks=None, + sparse_mask_fine=None, + only_qv=False, ) return o diff --git a/lightllm/models/vit/triton_kernel/flashattention_nopad.py b/lightllm/models/vit/triton_kernel/flashattention_nopad.py index b21464a02b..a9cbcda71d 100644 --- a/lightllm/models/vit/triton_kernel/flashattention_nopad.py +++ b/lightllm/models/vit/triton_kernel/flashattention_nopad.py @@ -203,6 +203,8 @@ def flash_attention_v3_fwd( pack_gqa=None, sm_margin=0, sinks=None, + sparse_mask_fine=None, + only_qv=False, ) return o diff --git a/unit_tests/models/vit/test_flash_attention_forward.py b/unit_tests/models/vit/test_flash_attention_forward.py index 53b70e6836..70849445b4 100644 --- a/unit_tests/models/vit/test_flash_attention_forward.py +++ b/unit_tests/models/vit/test_flash_attention_forward.py @@ -3,6 +3,16 @@ import time import pytest from lightllm.models.vit.triton_kernel.flashattention_nopad import flash_attention_fwd +from lightllm.models.vit.triton_kernel import flashattention_nopad +from lightllm.common.basemodel.attention_vit.fa3.fp import Fa3VitAttBackend + + +_fa3_mark = pytest.mark.skipif( + not flashattention_nopad._flash_attn_v3_available + or not torch.cuda.is_available() + or torch.cuda.get_device_capability()[0] < 9, + reason="FA3 requires sgl_kernel and a Hopper GPU", +) def reference_attention_varlen(q, k, v, cu): @@ -32,7 +42,21 @@ def reference_attention_varlen(q, k, v, cu): @pytest.mark.parametrize("dtype,atol", [(torch.float16, 1e-2), (torch.bfloat16, 2e-2)]) -def test_varlen(dtype, atol, batch=4, heads=8, d=80, device="cuda:0"): +@pytest.mark.parametrize( + "attention_fn", + [ + flash_attention_fwd, + pytest.param( + getattr(flashattention_nopad, "flash_attention_v3_fwd", None), + marks=_fa3_mark, + ), + pytest.param( + Fa3VitAttBackend._vit_att_fwd, + marks=_fa3_mark, + ), + ], +) +def test_varlen(dtype, atol, attention_fn, batch=4, heads=8, d=80, device="cuda:0"): torch.manual_seed(0) lengths = torch.randint(1, 257, (batch,)) max_len = int(lengths.max().item()) @@ -45,10 +69,10 @@ def test_varlen(dtype, atol, batch=4, heads=8, d=80, device="cuda:0"): k = torch.randn_like(q) v = torch.randn_like(q) out_tri = torch.randn_like(q) - flash_attention_fwd(q, k, v, out_tri, cu, max_len) + attention_fn(q, k, v, out_tri, cu, max_len) a = time.time() for _ in range(100): - flash_attention_fwd(q, k, v, out_tri, cu, max_len) + attention_fn(q, k, v, out_tri, cu, max_len) b = time.time() print(f"flash_attention_fwd time: {(b - a) / 100 * 1000:.2f} ms") out_ref = reference_attention_varlen(q, k, v, cu) From fb029ed4d9afb962d5422fdb8dfe254632a91bd4 Mon Sep 17 00:00:00 2001 From: sufubao Date: Thu, 15 Oct 2026 10:52:42 +0800 Subject: [PATCH 14/40] fix: separate draft tail graph layouts across padding --- lightllm/common/basemodel/basemodel.py | 47 ++--- lightllm/common/basemodel/cuda_graph.py | 6 +- .../proposers/eagle_with_att.py | 5 +- .../common/basemodel/test_draft_tail_graph.py | 167 ++++++++++++++++++ unit_tests/models/test_post_layer_output.py | 6 +- .../mtp_speculative/test_eagle_with_att.py | 5 +- 6 files changed, 206 insertions(+), 30 deletions(-) create mode 100644 unit_tests/common/basemodel/test_draft_tail_graph.py diff --git a/lightllm/common/basemodel/basemodel.py b/lightllm/common/basemodel/basemodel.py index 22c2cc4684..069179bb2c 100755 --- a/lightllm/common/basemodel/basemodel.py +++ b/lightllm/common/basemodel/basemodel.py @@ -277,7 +277,7 @@ def _align_decode_batch_size(self, batch_size: int) -> int: return triton.cdiv(batch_size, alignment) * alignment def use_replayssm_draft_tail(self): - """Keep first-draft graph shapes distinct from single-row draft decode.""" + """Project only accepted-tail rows in the first draft pass.""" return ( self.is_mtp_draft_model and self.args.ssm_state_mode == "replay" @@ -311,12 +311,21 @@ def _init_cudagraph(self): capture_infer_cost=self.args.mtp_dynamic_verify, ) ) + self.draft_tail_graph = None if self.graph is not None: if self.use_replayssm_draft_tail(): width = self.args.mtp_step + 1 - sizes = self.graph.cuda_graph_batch_sizes - self.graph.cuda_graph_batch_sizes = sorted(set(sizes + [width * n for n in sizes])) - self.graph.max_batch_size = self.graph.cuda_graph_batch_sizes[-1] + self.draft_tail_graph = CudaGraph( + batch_step_size_before_split=width, + split_batch_size=self.args.graph_split_batch_size * width, + batch_step_size_after_split=self.args.graph_grow_step_size * width, + max_batch_size=self.graph_max_batch_size * width, + max_len_in_batch=self.graph_max_len_in_batch, + tp_world_size=self.tp_world_size_, + ) + # Draft passes execute sequentially, so both layouts can share the pool. + self.draft_tail_graph.mempool = self.graph.mempool + self.draft_tail_graph.warmup(self, logits_row_stride=width) if get_env_start_args().enable_decode_microbatch_overlap: self.graph.warmup_overlap(self) else: @@ -374,22 +383,14 @@ def _create_inferstate(self, model_input: ModelInput, microbatch_index: int = 0) infer_state.is_draft_model = self.is_mtp_draft_model infer_state.return_all_prompt_logics = self.return_all_prompt_logics infer_state.batch_size = model_input.batch_size - # Above the logical request capacity, this can only be the first draft - # pass. Warmup and replay must use the same compact projection layout. - if ( - self.use_replayssm_draft_tail() - and not model_input.is_prefill - and model_input.batch_size > self.args.running_max_req_size - ): + if model_input.b_logits_rows is not None: + assert self.use_replayssm_draft_tail() and not model_input.is_prefill width = self.args.mtp_step + 1 assert model_input.batch_size % width == 0 count = model_input.batch_size // width rows = model_input.b_logits_rows - if rows is None: - rows = torch.arange(count, device=model_input.input_ids.device, dtype=torch.int64) * width + width - 1 - else: - assert rows.ndim == 1 and rows.numel() <= count - rows = F.pad(rows, (0, count - rows.numel())) + assert rows.ndim == 1 and rows.numel() <= count + rows = F.pad(rows, (0, count - rows.numel())) infer_state.b_logits_rows = rows infer_state.logits_row_stride = width infer_state.total_token_num = model_input.total_token_num @@ -638,20 +639,24 @@ def _decode( ) origin_batch_size = model_input.batch_size + graph = self.draft_tail_graph if model_input.b_logits_rows is not None else self.graph # 空 DP rank 也需要完整的 dummy verify group；TP/SP 切分不能破坏 MTP 分组。 infer_batch_size = self._align_decode_batch_size(max(1, origin_batch_size)) + if model_input.b_logits_rows is not None: + width = self.args.mtp_step + 1 + infer_batch_size = triton.cdiv(infer_batch_size, width) * width # CUDA Graph 可能继续向上对齐 batch size，并因此加入 seq_len=2 的 # dummy request。先用最终可能出现的 KV 长度判断 graph，再统一 padding 一次。 infer_max_kv_seq_len = max(2, model_input.max_kv_seq_len) - use_cuda_graph = self.graph is not None and self.graph.can_run( + use_cuda_graph = graph is not None and graph.can_run( batch_size=infer_batch_size, max_len_in_batch=infer_max_kv_seq_len, ) need_capture = False if use_cuda_graph: - infer_batch_size = self.graph.find_closest_graph_batch_size(batch_size=infer_batch_size) - need_capture = self.graph.need_capture(infer_batch_size) + infer_batch_size = graph.find_closest_graph_batch_size(batch_size=infer_batch_size) + need_capture = graph.need_capture(infer_batch_size) model_input = self._create_padded_decode_model_input(model_input=model_input, new_batch_size=infer_batch_size) infer_state = self._create_inferstate(model_input) @@ -663,9 +668,9 @@ def _decode( if use_cuda_graph: if need_capture: - model_output: ModelOutput = self.graph.capture_decode(self._token_forward, infer_state) + model_output: ModelOutput = graph.capture_decode(self._token_forward, infer_state) else: - model_output: ModelOutput = self.graph.replay(infer_state) + model_output: ModelOutput = graph.replay(infer_state) else: model_output = self._token_forward(infer_state) diff --git a/lightllm/common/basemodel/cuda_graph.py b/lightllm/common/basemodel/cuda_graph.py index ee888c1a00..fa236757be 100644 --- a/lightllm/common/basemodel/cuda_graph.py +++ b/lightllm/common/basemodel/cuda_graph.py @@ -250,7 +250,7 @@ def replay(self, infer_state, infer_state1=None): return self._replay(infer_state) @torch.no_grad() - def warmup(self, model): + def warmup(self, model, logits_row_stride=1): logger.info("Begin capture cudagraph, use the --disable_cudagraph to disable it.") # for typing easy from .basemodel import TpPartBaseModel @@ -286,6 +286,10 @@ def warmup(self, model): multimodal_params=[{"images": [], "audios": []} for _ in range(batch_size)], **model._gen_special_model_input(batch_size), ) + if logits_row_stride > 1: + model_input.b_logits_rows = torch.arange( + logits_row_stride - 1, batch_size, logits_row_stride, dtype=torch.int64, device="cuda" + ) model_output: ModelOutput = model.forward(model_input) del model_output del input_ids diff --git a/lightllm/server/router/model_infer/mtp_speculative/proposers/eagle_with_att.py b/lightllm/server/router/model_infer/mtp_speculative/proposers/eagle_with_att.py index 0f9c4fe8f1..d5dcdca05b 100644 --- a/lightllm/server/router/model_infer/mtp_speculative/proposers/eagle_with_att.py +++ b/lightllm/server/router/model_infer/mtp_speculative/proposers/eagle_with_att.py @@ -73,10 +73,7 @@ def propose_next( verify_draft_input = copy.copy(target_model_input) verify_draft_input.input_ids = target_next_token_ids verify_draft_input.mtp_draft_input_hiddens = target_model_output.mtp_collector.spec_hidden - if ( - draft_model.use_replayssm_draft_tail() - and verify_draft_input.batch_size > draft_model.args.running_max_req_size - ): + if draft_model.use_replayssm_draft_tail(): verify_draft_input.b_logits_rows = accepted_tail_rows extend_output = draft_model.forward(verify_draft_input) diff --git a/unit_tests/common/basemodel/test_draft_tail_graph.py b/unit_tests/common/basemodel/test_draft_tail_graph.py new file mode 100644 index 0000000000..5d14f988e8 --- /dev/null +++ b/unit_tests/common/basemodel/test_draft_tail_graph.py @@ -0,0 +1,167 @@ +from types import SimpleNamespace + +import pytest +import torch + +from lightllm.common.basemodel import basemodel, cuda_graph +from lightllm.common.basemodel.basemodel import TpPartBaseModel +from lightllm.common.basemodel.batch_objs import ModelInput, ModelMtpOutputCollector, ModelOutput +from lightllm.common.basemodel.hidden_collector import NoopHiddenCollector +from lightllm.common.basemodel.infer_struct import InferStateInfo +from lightllm.server.core.objs.start_args_type import StartArgs + + +@pytest.fixture +def draft_model(monkeypatch): + args = StartArgs( + ssm_state_mode="replay", + replayssm_projection_mode="parallel", + mtp_mode="eagle_with_att", + mtp_step=4, + running_max_req_size=264, + graph_max_batch_size=256, + max_req_total_len=128, + ) + monkeypatch.setattr(basemodel, "get_env_start_args", lambda: args) + monkeypatch.setattr(cuda_graph, "get_env_start_args", lambda: args) + monkeypatch.setattr(basemodel.dist_group_manager, "get_group", lambda _: None) + monkeypatch.setattr(InferStateInfo, "init_some_extra_state", lambda *args: None) + warmups = [] + monkeypatch.setattr(cuda_graph.CudaGraph, "warmup", lambda graph, model, **kw: warmups.append((graph, kw))) + + # Execute CPU row selection while retaining the real graph lookup, input + # copying, padding and output unpadding paths. + def capture(graph, forward, state): + output = forward(state) + + def replay(): + current = forward(state) + output.logits.copy_(current.logits) + output.mtp_collector.spec_hidden.copy_(current.mtp_collector.spec_hidden) + + graph.graph[state.batch_size] = (SimpleNamespace(replay=replay), state, output) + return output + + monkeypatch.setattr(cuda_graph.CudaGraph, "capture_decode", capture) + model = TpPartBaseModel.__new__(TpPartBaseModel) + model.args = args + model.is_mtp_draft_model = True + model.tp_world_size_ = 1 + model.config = {"model_type": "qwen3_5_text"} + model.disable_cudagraph = False + model.graph_max_batch_size = args.graph_max_batch_size + model.graph_max_len_in_batch = 128 + model.max_seq_length = 128 + model.mtp_manager = SimpleNamespace( + get_decode_tokens_per_request=lambda _: 1, get_decode_batch_alignment=lambda _: 1 + ) + model.infer_state_class = InferStateInfo + model.hidden_collector_prototype = NoopHiddenCollector() + model.return_all_prompt_logics = False + model.req_manager = SimpleNamespace(HOLD_REQUEST_ID=264) + model.mem_manager = None + model._select_mem_indexes = lambda _: None + model.decode_att_backend = SimpleNamespace( + create_att_decode_state=lambda **kw: SimpleNamespace( + init_state=lambda: None, copy_for_decode_cuda_graph=lambda _: None + ) + ) + model.decode_att_backend1 = None + + def forward(state): + hidden = state.input_ids[:, None].float() + logits = hidden if state.b_logits_rows is None else hidden.index_select(0, state.b_logits_rows) + return ModelOutput( + logits=logits, + logits_row_stride=state.logits_row_stride, + mtp_collector=ModelMtpOutputCollector(spec_hidden=hidden), + ) + + model._token_forward = forward + model._init_cudagraph() + assert warmups == [(model.draft_tail_graph, {"logits_row_stride": 5}), (model.graph, {})] + return model + + +def make_input(rows, accepted=None): + return ModelInput( + batch_size=rows, + total_token_num=rows * 10, + max_q_seq_len=1, + max_kv_seq_len=10, + input_ids=torch.arange(rows), + b_req_idx=torch.arange(rows, dtype=torch.int32) // (5 if accepted is not None else 1), + b_mtp_index=torch.zeros(rows, dtype=torch.int32), + b_seq_len=torch.full((rows,), 10, dtype=torch.int32), + b_position_delta=torch.zeros(rows, dtype=torch.int32), + b_shared_seq_len=torch.zeros(rows, dtype=torch.int32), + b_shared_radix_node_id=torch.zeros(rows, dtype=torch.int64), + multimodal_params=[{} for _ in range(rows)], + b_logits_rows=accepted, + ) + + +@pytest.mark.parametrize("eager", [False, True]) +@pytest.mark.parametrize("requests", [0, 3, 32, 51, 52, 53, 264]) +def test_first_draft_selects_accepted_tail_across_padding(draft_model, eager, requests): + model = draft_model + if eager: + model.graph = model.draft_tail_graph = None + for step in range(2): + accepted = torch.arange(requests) * 5 + (torch.arange(requests) + step) % 5 + model_input = make_input(requests * 5, accepted) + output = model._decode(model_input) + torch.testing.assert_close(output.logits[:, 0], accepted.float()) + torch.testing.assert_close(output.mtp_collector.spec_hidden[:, 0], model_input.input_ids.float()) + assert output.logits_row_stride == 5 + assert model_input.batch_size == requests * 5 + assert model_input.b_logits_rows is accepted + if not eager and requests == 52: + assert 320 in model.draft_tail_graph.graph + assert not model.graph.graph + + +@pytest.mark.parametrize("rows", [0, 160, 256, 260, 263, 264]) +def test_later_draft_keeps_one_output_per_request(draft_model, rows): + model = draft_model + # 160 physical rows occur in both layouts; their captured outputs must not mix. + model._decode(make_input(160, torch.arange(32) * 5 + 1)) + for step in range(2): + model_input = make_input(rows) + model_input.input_ids += step + output = model._decode(model_input) + assert output.logits_row_stride == 1 + torch.testing.assert_close(output.logits[:, 0], model_input.input_ids.float()) + assert output.mtp_collector.spec_hidden.shape == (rows, 1) + if rows == 160: + assert 160 in model.graph.graph and 160 in model.draft_tail_graph.graph + assert model.graph.graph[160][2].logits.shape == (160, 1) + assert model.draft_tail_graph.graph[160][2].logits.shape == (32, 1) + elif rows > 256: + assert not model.graph.graph + + +def test_tail_graph_warmup_marks_selected_rows(monkeypatch): + args = StartArgs() + monkeypatch.setattr(cuda_graph, "get_env_start_args", lambda: args) + + def on_cpu(factory): + def allocate(*args, **kwargs): + kwargs["device"] = "cpu" + return factory(*args, **kwargs) + + return allocate + + for name in ("tensor", "full", "zeros", "arange"): + monkeypatch.setattr(torch, name, on_cpu(getattr(torch, name))) + seen = [] + model = SimpleNamespace( + req_manager=SimpleNamespace(HOLD_REQUEST_ID=8, free_all=lambda: None), + mem_manager=SimpleNamespace(free_all=lambda: None), + _gen_special_model_input=lambda _: {}, + forward=lambda model_input: seen.append(model_input), + ) + graph = cuda_graph.CudaGraph(5, 5, 5, max_batch_size=10) + graph.warmup(model, logits_row_stride=5) + assert [x.batch_size for x in seen] == [10, 5] + assert [x.b_logits_rows.tolist() for x in seen] == [[4, 9], [4]] diff --git a/unit_tests/models/test_post_layer_output.py b/unit_tests/models/test_post_layer_output.py index e21d0b1d1c..c8e081862d 100644 --- a/unit_tests/models/test_post_layer_output.py +++ b/unit_tests/models/test_post_layer_output.py @@ -60,7 +60,7 @@ def candidates(*args): args = SimpleNamespace(target_vocab_topk_sampling=None, draft_vocab_topk_sampling=None) setattr(args, "draft_vocab_topk_sampling" if draft else "target_vocab_topk_sampling", 16 if enabled else None) monkeypatch.setattr(llama_post, "get_env_start_args", lambda: args) - state = SimpleNamespace(is_draft_model=draft, dist_group=None) + state = SimpleNamespace(is_draft_model=draft, dist_group=None, b_logits_rows=None) weight = SimpleNamespace(lm_head_weight_=Weight()) method = head._draft_lm_head_and_gather if draft else head._target_lm_head_and_gather output = method(torch.ones(3, 2), 3, weight, state) @@ -290,7 +290,9 @@ def test_model_output_and_unpadding_preserve_candidate_mapping() -> None: model = TpPartBaseModel.__new__(TpPartBaseModel) logits = torch.randn(3, 2) ids = torch.tensor([[10, 20], [30, 40], [50, 60]]) - state = SimpleNamespace(hidden_collector=NoopHiddenCollector(), prompt_logics=None, ssm_run_config=None) + state = SimpleNamespace( + hidden_collector=NoopHiddenCollector(), prompt_logics=None, ssm_run_config=None, logits_row_stride=1 + ) output = model._create_model_output(PostLayerOutput(logits, ids), state) unpadded = model._create_unpad_decode_model_output(output, 2) torch.testing.assert_close(unpadded.logits, logits[:2]) diff --git a/unit_tests/server/router/model_infer/mtp_speculative/test_eagle_with_att.py b/unit_tests/server/router/model_infer/mtp_speculative/test_eagle_with_att.py index 894833ab03..94be2ef52a 100644 --- a/unit_tests/server/router/model_infer/mtp_speculative/test_eagle_with_att.py +++ b/unit_tests/server/router/model_infer/mtp_speculative/test_eagle_with_att.py @@ -216,7 +216,8 @@ def forward(model_input): assert target_input.mtp_draft_input_hiddens is None -def test_selected_tail_logits_keep_complete_draft_kv_input(): +@pytest.mark.parametrize("capacity", [2, 10, 264]) +def test_selected_tail_logits_keep_complete_draft_kv_input(capacity): hidden = torch.arange(20).view(10, 2).float() target_input = SimpleNamespace( is_prefill=False, @@ -237,7 +238,7 @@ def forward(model_input): draft_model = SimpleNamespace( forward=forward, use_replayssm_draft_tail=lambda: True, - args=SimpleNamespace(running_max_req_size=2), + args=SimpleNamespace(running_max_req_size=capacity), ) proposer = EagleWithAttProposer( backend=SimpleNamespace(draft_models=[draft_model], _gen_argmax_token_ids=lambda out: out.logits[:, 0].long()), From 16d55f81358d8663e564df199ca5c924bf627a67 Mon Sep 17 00:00:00 2001 From: sufubao Date: Thu, 15 Oct 2026 11:17:14 +0800 Subject: [PATCH 15/40] test: cover continuous ReplaySSM derived state lifecycle --- ...=4,varlen=true}_NVIDIA_H100_80GB_HBM3.json | 12 --- .../linear_att/test_replayssm_parallel.py | 92 +++++++++++++++++++ 2 files changed, 92 insertions(+), 12 deletions(-) delete mode 100644 lightllm/common/triton_utils/autotune_kernel_configs/triton_3.6.0/NVIDIA_H100_80GB_HBM3/ssm_update_cycle:v1/{dtype=(torch.bfloat16,torch.bfloat16),gate=0,history=8,history_dtype=torch.bfloat16,mode=replay,projection=block_mma_activation_precision,shape=(48,4,12,4,128,128,4),strides=(4120,4120,4120,4120,4120),v=4,varlen=true}_NVIDIA_H100_80GB_HBM3.json diff --git a/lightllm/common/triton_utils/autotune_kernel_configs/triton_3.6.0/NVIDIA_H100_80GB_HBM3/ssm_update_cycle:v1/{dtype=(torch.bfloat16,torch.bfloat16),gate=0,history=8,history_dtype=torch.bfloat16,mode=replay,projection=block_mma_activation_precision,shape=(48,4,12,4,128,128,4),strides=(4120,4120,4120,4120,4120),v=4,varlen=true}_NVIDIA_H100_80GB_HBM3.json b/lightllm/common/triton_utils/autotune_kernel_configs/triton_3.6.0/NVIDIA_H100_80GB_HBM3/ssm_update_cycle:v1/{dtype=(torch.bfloat16,torch.bfloat16),gate=0,history=8,history_dtype=torch.bfloat16,mode=replay,projection=block_mma_activation_precision,shape=(48,4,12,4,128,128,4),strides=(4120,4120,4120,4120,4120),v=4,varlen=true}_NVIDIA_H100_80GB_HBM3.json deleted file mode 100644 index 71971ee8d7..0000000000 --- a/lightllm/common/triton_utils/autotune_kernel_configs/triton_3.6.0/NVIDIA_H100_80GB_HBM3/ssm_update_cycle:v1/{dtype=(torch.bfloat16,torch.bfloat16),gate=0,history=8,history_dtype=torch.bfloat16,mode=replay,projection=block_mma_activation_precision,shape=(48,4,12,4,128,128,4),strides=(4120,4120,4120,4120,4120),v=4,varlen=true}_NVIDIA_H100_80GB_HBM3.json +++ /dev/null @@ -1,12 +0,0 @@ -{ - "192": { - "BV": 32, - "num_stages": 3, - "num_warps": 1 - }, - "128": { - "BV": 32, - "num_stages": 3, - "num_warps": 1 - } -} diff --git a/unit_tests/common/basemodel/triton_kernel/linear_att/test_replayssm_parallel.py b/unit_tests/common/basemodel/triton_kernel/linear_att/test_replayssm_parallel.py index 1eafafae20..0a2248080c 100644 --- a/unit_tests/common/basemodel/triton_kernel/linear_att/test_replayssm_parallel.py +++ b/unit_tests/common/basemodel/triton_kernel/linear_att/test_replayssm_parallel.py @@ -69,3 +69,95 @@ def test_parallel_raw_against_sequential_reference(width): projection_mode="parallel", run_config={"BV": 128, "num_warps": 2, "num_stages": 1}, ) + + +@pytest.mark.parametrize("capacity,width", [(8, 3), (16, 5), (64, 16)]) +@pytest.mark.parametrize("capture", [False, True], ids=["eager", "graph"]) +def test_parallel_derived_continuous_state(capacity, width, capture): + from unit_tests.common.basemodel.triton_kernel.linear_att.test_replayssm import reference + + torch.manual_seed(731) + state = torch.randn(2, 4, 4, 128, 128, device="cuda", dtype=torch.bfloat16) * 0.02 + initial = state.clone() + cache = ReplaySSMCache(state, capacity, width, num_key_heads=2, projection_mode="parallel", fold_method="derived") + for history in (cache.keys, cache.deltas, cache.gates, cache.raw_keys, cache.raw_values, cache.betas): + history.fill_(float("nan")) + reqs = torch.tensor([2, 0, 3], device="cuda", dtype=torch.int32) + lengths = [width, width - 1] + cu = torch.tensor([0, width, 2 * width - 1, 2 * width], device="cuda", dtype=torch.int32) + q = torch.randn(1, 2 * width, 2, 128, device="cuda", dtype=torch.bfloat16) + k = torch.randn_like(q) + v = torch.randn(1, 2 * width, 4, 128, device="cuda", dtype=torch.bfloat16) + a = torch.randn(2 * width, 4, device="cuda", dtype=torch.bfloat16) - 3 + b = torch.randn_like(a) + alog = torch.randn(4, device="cuda") * 0.1 + bias = torch.randn_like(alog) * 0.1 + accepted = torch.zeros(4, device="cuda", dtype=torch.int32) + + def step(): + cache.prepare_decode(reqs, cu) + outputs = [cache.forward(layer, q, k, v, a, b, alog, bias, reqs, cu) for layer in range(2)] + cache.accept_updates(reqs, accepted) + return outputs + + # Compile before capture, then reset the live state and capture only HOLD rows. + step() + state.copy_(initial) + cache.cursors.zero_() + if capture: + reqs.fill_(cache.hold) + torch.cuda.synchronize() + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + outputs = step() + reqs.copy_(torch.tensor([2, 0, 3], device="cuda", dtype=torch.int32)) + + expected = initial.float() + pending = [0, 0] + folds = [0, 0] + for iteration in range(3 * capacity + 5): + for tensor in (q, k, v, b): + tensor.normal_() + a.normal_(-3, 1) + counts = [iteration % width + 1, (iteration + 1) % (width - 1) + 1] + for seq, req in enumerate((2, 0)): + accepted[req] = counts[seq] - 1 + if pending[seq] + lengths[seq] > capacity: + expected[:, req] = expected[:, req].bfloat16().float() + pending[seq] = 0 + folds[seq] += 1 + pending[seq] += counts[seq] + if capture: + graph.replay() + else: + outputs = step() + for layer in range(2): + for seq, req in enumerate((2, 0)): + current = expected[layer, req].clone() + for j in range(lengths[seq]): + token = seq * width + j + output, current = reference( + q[0, token], k[0, token], v[0, token], a[token], b[token], alog, bias, current + ) + torch.testing.assert_close(outputs[layer][0, token].float(), output, atol=0.006, rtol=0.02) + if j + 1 == counts[seq]: + expected[layer, req].copy_(current) + assert torch.count_nonzero(outputs[layer][0, -1]) == 0 + + if iteration == capacity: + saved_state, saved_cursors = state.clone(), cache.cursors.clone() + snapshot = cache.snapshot_accepted_state(2) + torch.testing.assert_close(snapshot.float(), expected[:, 2], atol=0.002, rtol=0.02) + torch.testing.assert_close(state, saved_state, atol=0, rtol=0) + torch.testing.assert_close(cache.cursors, saved_cursors, atol=0, rtol=0) + # Restore the canonical checkpoint while another request keeps its history. + cache.clear_history(2) + state[:, 2].copy_(snapshot) + expected[:, 2].copy_(snapshot.float()) + pending[0] = 0 + + assert min(folds) >= 2 # Exercise both ring halves, including after restore. + cache.merge_accepted_updates(reqs[:2]) + torch.testing.assert_close(state.float(), expected.bfloat16().float(), atol=0.002, rtol=0.02) + torch.testing.assert_close(state[:, [1, 3]], initial[:, [1, 3]], atol=0, rtol=0) + assert torch.count_nonzero(cache.cursors) == 0 From 7dcda704e93ac56eac904fd47b1ad72b2d8a8c0a Mon Sep 17 00:00:00 2001 From: sufubao Date: Thu, 15 Oct 2026 14:32:21 +0800 Subject: [PATCH 16/40] refactor: group ReplaySSM options in a JSON config --- lightllm/common/basemodel/basemodel.py | 2 +- lightllm/common/req_manager/linear_att.py | 6 +- lightllm/server/api_cli.py | 25 ++----- lightllm/server/api_start.py | 45 +++++++---- lightllm/server/core/objs/start_args_type.py | 4 +- .../common/basemodel/test_draft_tail_graph.py | 2 +- .../common/basemodel/test_model_output.py | 4 +- unit_tests/server/test_ssm_state_mode.py | 75 ++++++++++++++++--- 8 files changed, 110 insertions(+), 53 deletions(-) diff --git a/lightllm/common/basemodel/basemodel.py b/lightllm/common/basemodel/basemodel.py index 069179bb2c..d328fbd501 100755 --- a/lightllm/common/basemodel/basemodel.py +++ b/lightllm/common/basemodel/basemodel.py @@ -281,7 +281,7 @@ def use_replayssm_draft_tail(self): return ( self.is_mtp_draft_model and self.args.ssm_state_mode == "replay" - and self.args.replayssm_projection_mode == "parallel" + and self.args.ssm_replay_config.get("compute_mode", "inline") == "parallel" and self.args.mtp_mode == "eagle_with_att" and self.args.mtp_step > 0 and not self.args.mtp_dynamic_verify diff --git a/lightllm/common/req_manager/linear_att.py b/lightllm/common/req_manager/linear_att.py index 020c64ae9e..6a96e2ac22 100644 --- a/lightllm/common/req_manager/linear_att.py +++ b/lightllm/common/req_manager/linear_att.py @@ -62,13 +62,13 @@ def __init__( self.ssm_update_cache = ReplaySSMCache( self.req_to_ssm_state.buffer, - args.replayssm_cache_len, + args.ssm_replay_config.get("cache_len", 16), self.mtp_step + 1, linear_config.conv_state_dtype, num_key_heads=linear_config.num_linear_k_heads, kda=self.kda, - projection_mode=args.replayssm_projection_mode, - fold_method=args.replayssm_fold_method, + projection_mode=args.ssm_replay_config.get("compute_mode", "inline"), + fold_method=args.ssm_replay_config.get("fold_method", "raw"), ) elif args.ssm_state_mode == "compact": from lightllm.common.basemodel.triton_kernel.linear_att.replayssm_compact import CompactSSMCache diff --git a/lightllm/server/api_cli.py b/lightllm/server/api_cli.py index 2c9309d641..9cb79fe85e 100644 --- a/lightllm/server/api_cli.py +++ b/lightllm/server/api_cli.py @@ -1,4 +1,5 @@ import argparse +import json def add_cli_args(parser: argparse.ArgumentParser) -> argparse.ArgumentParser: @@ -955,24 +956,12 @@ def add_cli_args(parser: argparse.ArgumentParser) -> argparse.ArgumentParser: help="SSM state storage: native snapshots, compact MTP records, or deferred GDN/KDA replay (FP32/BF16).", ) parser.add_argument( - "--replayssm_cache_len", - type=int, - default=16, - choices=[4, 8, 16, 32, 64], - help="History capacity for ssm_state_mode=replay only; must cover the MTP verify width.", - ) - parser.add_argument( - "--replayssm_projection_mode", - choices=["inline", "precompute", "parallel"], - default="inline", - help="ReplaySSM projection; parallel uses a Tensor Core GDN MTP solver (BF16, 128x128 heads).", - ) - parser.add_argument( - "--replayssm_fold_method", - choices=["raw", "derived"], - default="raw", - help="ReplaySSM fold: raw replays accepted inputs; derived reuses parallel-solver updates and changes " - "checkpoint rounding. Derived requires parallel projection; state precision and capacity are unchanged.", + "--ssm_replay_config", + type=json.loads, + default={}, + help='Replay-only JSON object: {"cache_len":16,"compute_mode":"inline","fold_method":"raw"}. ' + "cache_len: 4/8/16/32/64, at least the verify width; compute_mode: inline/precompute/parallel; " + "fold_method: raw/derived. Parallel requires BF16 GDN MTP with 128x128 heads; derived requires parallel.", ) parser.add_argument( "--linear_att_ssm_data_type", diff --git a/lightllm/server/api_start.py b/lightllm/server/api_start.py index e76a92e1ff..8b68306fd4 100644 --- a/lightllm/server/api_start.py +++ b/lightllm/server/api_start.py @@ -36,17 +36,33 @@ def _set_envs_and_config(args: StartArgs): def _validate_ssm_state_mode(args: StartArgs): assert args.ssm_state_mode in ("native", "compact", "replay"), "Invalid ssm_state_mode" - assert args.replayssm_projection_mode in ("inline", "precompute", "parallel"), "Invalid replayssm_projection_mode" - assert args.replayssm_fold_method in ("raw", "derived"), "Invalid replayssm_fold_method" - assert ( - args.replayssm_fold_method == "raw" or args.replayssm_projection_mode == "parallel" - ), "derived fold requires parallel ReplaySSM projection" - assert ( - args.replayssm_projection_mode == "inline" or args.ssm_state_mode == "replay" - ), "Non-inline ReplaySSM projection requires ssm_state_mode=replay" - if args.replayssm_projection_mode == "parallel": - assert args.mtp_step > 0, "Parallel ReplaySSM requires MTP" - assert args.linear_att_ssm_data_type == "bfloat16", "Parallel ReplaySSM requires BF16 state" + config = args.ssm_replay_config + assert isinstance(config, dict), "ssm_replay_config must be a JSON object" + assert args.ssm_state_mode == "replay" or not config, "ssm_replay_config requires ssm_state_mode=replay" + if args.ssm_state_mode == "replay": + defaults = {"cache_len": 16, "compute_mode": "inline", "fold_method": "raw"} + unknown = config.keys() - defaults.keys() + assert not unknown, f"Unknown ssm_replay_config fields: {unknown}" + config = {**defaults, **config} + assert type(config["cache_len"]) is int and config["cache_len"] in ( + 4, + 8, + 16, + 32, + 64, + ), "ReplaySSM capacity must be a power of two and at least 4, up to 64" + assert config["compute_mode"] in ("inline", "precompute", "parallel"), "Invalid replay compute_mode" + assert config["fold_method"] in ("raw", "derived"), "Invalid replay fold_method" + assert ( + config["fold_method"] == "raw" or config["compute_mode"] == "parallel" + ), "derived fold requires parallel ReplaySSM compute_mode" + if config["compute_mode"] == "parallel": + assert args.mtp_step > 0, "Parallel ReplaySSM requires MTP" + assert args.linear_att_ssm_data_type == "bfloat16", "Parallel ReplaySSM requires BF16 state" + args.ssm_replay_config = config + logger.info(f"SSM mode: replay, config: {config}") + else: + logger.info(f"SSM mode: {args.ssm_state_mode}") if args.ssm_state_mode == "native": return assert args.linear_att_ssm_data_type in ( @@ -68,11 +84,8 @@ def _validate_ssm_state_mode(args: StartArgs): assert args.mtp_step > 0, "ssm_state_mode=compact requires MTP (mtp_step > 0)" else: assert args.mtp_step >= 0, "ssm_state_mode=replay requires mtp_step >= 0" - capacity = args.replayssm_cache_len - assert ( - capacity >= 4 and capacity & (capacity - 1) == 0 - ), "ReplaySSM capacity must be a power of two and at least 4" - assert args.replayssm_cache_len >= args.mtp_step + 1, "ReplaySSM capacity must cover the verify width" + capacity = config["cache_len"] + assert capacity >= args.mtp_step + 1, "ReplaySSM capacity must cover the verify width" def _launch_subprocesses(args: StartArgs): diff --git a/lightllm/server/core/objs/start_args_type.py b/lightllm/server/core/objs/start_args_type.py index f0f7792450..7ccaf7d78f 100644 --- a/lightllm/server/core/objs/start_args_type.py +++ b/lightllm/server/core/objs/start_args_type.py @@ -245,9 +245,7 @@ class StartArgs: # hybrid attention model (Qwen3Next) ssm_state_mode: str = "native" - replayssm_cache_len: int = 16 - replayssm_projection_mode: str = "inline" - replayssm_fold_method: str = "raw" + ssm_replay_config: dict = field(default_factory=dict) 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/unit_tests/common/basemodel/test_draft_tail_graph.py b/unit_tests/common/basemodel/test_draft_tail_graph.py index 5d14f988e8..e1aa933a4a 100644 --- a/unit_tests/common/basemodel/test_draft_tail_graph.py +++ b/unit_tests/common/basemodel/test_draft_tail_graph.py @@ -15,7 +15,7 @@ def draft_model(monkeypatch): args = StartArgs( ssm_state_mode="replay", - replayssm_projection_mode="parallel", + ssm_replay_config={"compute_mode": "parallel"}, mtp_mode="eagle_with_att", mtp_step=4, running_max_req_size=264, diff --git a/unit_tests/common/basemodel/test_model_output.py b/unit_tests/common/basemodel/test_model_output.py index 95acb17060..7c0606ac63 100644 --- a/unit_tests/common/basemodel/test_model_output.py +++ b/unit_tests/common/basemodel/test_model_output.py @@ -254,7 +254,7 @@ def test_draft_tail_scope(disabled): model = TpPartBaseModel.__new__(TpPartBaseModel) model.args = StartArgs( - ssm_state_mode="replay", replayssm_projection_mode="parallel", mtp_mode="eagle_with_att", mtp_step=4 + ssm_state_mode="replay", ssm_replay_config={"compute_mode": "parallel"}, mtp_mode="eagle_with_att", mtp_step=4 ) model.is_mtp_draft_model = disabled != "target" model.tp_world_size_ = 2 if disabled == "tp" else 1 @@ -262,5 +262,5 @@ def test_draft_tail_scope(disabled): model.args.mtp_dynamic_verify = disabled == "dynamic" model.args.enable_decode_microbatch_overlap = disabled == "overlap" if disabled == "projection": - model.args.replayssm_projection_mode = "inline" + model.args.ssm_replay_config = {"compute_mode": "inline"} assert model.use_replayssm_draft_tail() == (disabled is None) diff --git a/unit_tests/server/test_ssm_state_mode.py b/unit_tests/server/test_ssm_state_mode.py index 6016cd2a31..4675945059 100644 --- a/unit_tests/server/test_ssm_state_mode.py +++ b/unit_tests/server/test_ssm_state_mode.py @@ -1,3 +1,4 @@ +import json import pytest from lightllm.server.api_cli import make_argument_parser @@ -36,8 +37,8 @@ def test_ssm_mode_validation(monkeypatch, mode, dtype, step, model, capacity, er dtype, "--mtp_step", str(step), - "--replayssm_cache_len", - str(capacity), + "--ssm_replay_config", + json.dumps({"cache_len": capacity} if mode == "replay" else {}), ] ) if error: @@ -50,7 +51,7 @@ def test_ssm_mode_validation(monkeypatch, mode, dtype, step, model, capacity, er def test_ssm_mode_defaults_and_removed_flag(): parser = make_argument_parser() assert parser.parse_args([]).ssm_state_mode == StartArgs().ssm_state_mode == "native" - assert parser.parse_args([]).replayssm_projection_mode == StartArgs().replayssm_projection_mode == "inline" + assert parser.parse_args([]).ssm_replay_config == StartArgs().ssm_replay_config == {} with pytest.raises(SystemExit): parser.parse_args(["--enable_replayssm"]) with pytest.raises(SystemExit): @@ -60,8 +61,8 @@ def test_ssm_mode_defaults_and_removed_flag(): @pytest.mark.parametrize( "overrides,error", [ - ({"replayssm_cache_len": 2}, "power of two and at least 4"), - ({"replayssm_cache_len": 6}, "power of two and at least 4"), + ({"ssm_replay_config": {"cache_len": 2}}, "power of two and at least 4"), + ({"ssm_replay_config": {"cache_len": 6}}, "power of two and at least 4"), ({"mtp_step": -1}, "mtp_step >= 0"), ({"linear_att_ssm_data_type": "float16"}, "FP32 or BF16"), ], @@ -78,7 +79,7 @@ def test_programmatic_replay_config_rejected_before_model_loading(monkeypatch, o def test_replay_projection_mode_requires_replay(monkeypatch): monkeypatch.setattr("lightllm.utils.config_utils.get_model_type", lambda _: "qwen3_5") parser = make_argument_parser() - args = parser.parse_args(["--ssm_state_mode", "replay", "--replayssm_projection_mode", "precompute"]) + args = parser.parse_args(["--ssm_state_mode", "replay", "--ssm_replay_config", '{"compute_mode":"precompute"}']) _validate_ssm_state_mode(args) args.ssm_state_mode = "compact" with pytest.raises(AssertionError, match="requires ssm_state_mode=replay"): @@ -92,8 +93,7 @@ def test_parallel_replay_configuration(monkeypatch, fold): monkeypatch.setattr("lightllm.utils.config_utils.get_model_type", lambda _: "qwen3_5") args = StartArgs( ssm_state_mode="replay", - replayssm_projection_mode="parallel", - replayssm_fold_method=fold, + ssm_replay_config={"compute_mode": "parallel", "fold_method": fold}, mtp_step=4, mtp_mode="eagle_with_att", linear_att_ssm_data_type="bfloat16", @@ -107,6 +107,63 @@ def test_parallel_replay_configuration(monkeypatch, fold): def test_derived_replay_requires_parallel_configuration(): from lightllm.server.core.objs.start_args_type import StartArgs - args = StartArgs(ssm_state_mode="replay", replayssm_fold_method="derived") + args = StartArgs(ssm_state_mode="replay", ssm_replay_config={"fold_method": "derived"}) with pytest.raises(AssertionError, match="requires parallel"): _validate_ssm_state_mode(args) + + +@pytest.mark.parametrize( + "config,error", + [ + ([], "JSON object"), + (None, "JSON object"), + ({"projection_mode": "inline"}, "Unknown"), + ({"cache_len": True}, "capacity"), + ({"cache_len": 16.0}, "capacity"), + ({"cache_len": "16"}, "capacity"), + ({"cache_len": 128}, "capacity"), + ({"compute_mode": "auto"}, "compute_mode"), + ({"fold_method": "auto"}, "fold_method"), + ], +) +def test_invalid_replay_config(config, error): + args = StartArgs(ssm_state_mode="replay", ssm_replay_config=config) + with pytest.raises(AssertionError, match=error): + _validate_ssm_state_mode(args) + + +@pytest.mark.parametrize("mode", ["native", "compact"]) +def test_replay_config_rejects_other_modes(mode): + args = StartArgs(ssm_state_mode=mode, ssm_replay_config={"cache_len": 16}) + with pytest.raises(AssertionError, match="requires ssm_state_mode=replay"): + _validate_ssm_state_mode(args) + + +def test_partial_replay_config_defaults_and_serialization(monkeypatch): + from dataclasses import asdict + + monkeypatch.setattr("lightllm.utils.config_utils.get_model_type", lambda _: "qwen3_5") + args = make_argument_parser().parse_args( + ["--ssm_state_mode", "replay", "--ssm_replay_config", '{"compute_mode":"precompute"}'] + ) + _validate_ssm_state_mode(args) + assert args.ssm_replay_config == {"cache_len": 16, "compute_mode": "precompute", "fold_method": "raw"} + original = StartArgs(ssm_state_mode="replay", ssm_replay_config=args.ssm_replay_config) + restored = StartArgs(**json.loads(json.dumps(asdict(original)))) + _validate_ssm_state_mode(restored) + assert restored.ssm_replay_config == args.ssm_replay_config + assert StartArgs().ssm_replay_config == {} + + +@pytest.mark.parametrize( + "option,value", + [ + ("--replayssm_cache_len", "16"), + ("--replayssm_projection_mode", "inline"), + ("--replayssm_fold_method", "raw"), + ("--ssm_replay_config", "{invalid}"), + ], +) +def test_removed_flags_and_malformed_json(option, value): + with pytest.raises(SystemExit): + make_argument_parser().parse_args([option, value]) From 833d5eb05bbaac4944a7e8d2e0ea02699da37e59 Mon Sep 17 00:00:00 2001 From: sufubao Date: Thu, 15 Oct 2026 14:51:44 +0800 Subject: [PATCH 17/40] fix: reject unsupported compact MTP widths at startup --- lightllm/server/api_start.py | 4 ++++ unit_tests/server/test_ssm_state_mode.py | 6 +++++- 2 files changed, 9 insertions(+), 1 deletion(-) diff --git a/lightllm/server/api_start.py b/lightllm/server/api_start.py index 8b68306fd4..63417317ff 100644 --- a/lightllm/server/api_start.py +++ b/lightllm/server/api_start.py @@ -82,6 +82,10 @@ def _validate_ssm_state_mode(args: StartArgs): ), "Non-native SSM state modes currently require a GDN or KDA model" if args.ssm_state_mode == "compact": assert args.mtp_step > 0, "ssm_state_mode=compact requires MTP (mtp_step > 0)" + assert args.mtp_step <= 15, ( + "ssm_state_mode=compact requires mtp_step <= 15: " + "linear attention MTP state acceptance supports verify width <= 16" + ) else: assert args.mtp_step >= 0, "ssm_state_mode=replay requires mtp_step >= 0" capacity = config["cache_len"] diff --git a/unit_tests/server/test_ssm_state_mode.py b/unit_tests/server/test_ssm_state_mode.py index 4675945059..fb6591be51 100644 --- a/unit_tests/server/test_ssm_state_mode.py +++ b/unit_tests/server/test_ssm_state_mode.py @@ -12,7 +12,11 @@ ("native", "bfloat16", 0, "llama", 16, None), ("compact", "bfloat16", 2, "qwen3_5", 16, None), ("compact", "float32", 2, "qwen3_5", 16, None), - ("compact", "float32", 32, "qwen3_5", 16, None), + ("compact", "float32", 15, "qwen3_5", 16, None), + ("compact", "float32", 16, "qwen3_5", 16, "mtp_step <= 15"), + ("compact", "float32", 32, "qwen3_5", 16, "mtp_step <= 15"), + ("compact", "bfloat16", 15, "glm5_next", 16, None), + ("compact", "bfloat16", 16, "glm5_next", 16, "mtp_step <= 15"), ("compact", "bfloat16", 2, "glm5_next", 16, None), ("replay", "float32", 0, "glm5_next", 8, None), ("replay", "bfloat16", 2, "glm5_next_text", 8, None), From aca22817ddd2105a3d263decef5fa70f63f53424 Mon Sep 17 00:00:00 2001 From: sufubao Date: Thu, 15 Oct 2026 15:09:53 +0800 Subject: [PATCH 18/40] fix: validate GDN and KDA MTP widths across state modes --- lightllm/server/api_start.py | 25 ++++++++++--------- unit_tests/server/test_ssm_state_mode.py | 31 +++++++++++++++++++++++- 2 files changed, 43 insertions(+), 13 deletions(-) diff --git a/lightllm/server/api_start.py b/lightllm/server/api_start.py index 63417317ff..24ebf38ea5 100644 --- a/lightllm/server/api_start.py +++ b/lightllm/server/api_start.py @@ -63,15 +63,9 @@ def _validate_ssm_state_mode(args: StartArgs): logger.info(f"SSM mode: replay, config: {config}") else: logger.info(f"SSM mode: {args.ssm_state_mode}") - if args.ssm_state_mode == "native": - return - assert args.linear_att_ssm_data_type in ( - "float32", - "bfloat16", - ), "Non-native SSM state modes require FP32 or BF16 state" from lightllm.utils.config_utils import get_model_type - assert get_model_type(args.model_dir) in ( + is_gdn_or_kda = get_model_type(args.model_dir) in ( "glm5_next", "glm5_next_text", "qwen3_next", @@ -79,13 +73,20 @@ def _validate_ssm_state_mode(args: StartArgs): "qwen3_5_moe", "qwen3_5_text", "qwen3_5_moe_text", - ), "Non-native SSM state modes currently require a GDN or KDA model" - if args.ssm_state_mode == "compact": - assert args.mtp_step > 0, "ssm_state_mode=compact requires MTP (mtp_step > 0)" + ) + if is_gdn_or_kda: assert args.mtp_step <= 15, ( - "ssm_state_mode=compact requires mtp_step <= 15: " - "linear attention MTP state acceptance supports verify width <= 16" + "GDN/KDA requires mtp_step <= 15: " "linear attention MTP state acceptance supports verify width <= 16" ) + if args.ssm_state_mode == "native": + return + assert args.linear_att_ssm_data_type in ( + "float32", + "bfloat16", + ), "Non-native SSM state modes require FP32 or BF16 state" + assert is_gdn_or_kda, "Non-native SSM state modes currently require a GDN or KDA model" + if args.ssm_state_mode == "compact": + assert args.mtp_step > 0, "ssm_state_mode=compact requires MTP (mtp_step > 0)" else: assert args.mtp_step >= 0, "ssm_state_mode=replay requires mtp_step >= 0" capacity = config["cache_len"] diff --git a/unit_tests/server/test_ssm_state_mode.py b/unit_tests/server/test_ssm_state_mode.py index fb6591be51..7f1aad5e74 100644 --- a/unit_tests/server/test_ssm_state_mode.py +++ b/unit_tests/server/test_ssm_state_mode.py @@ -10,6 +10,7 @@ "mode,dtype,step,model,capacity,error", [ ("native", "bfloat16", 0, "llama", 16, None), + ("native", "bfloat16", 32, "llama", 16, None), ("compact", "bfloat16", 2, "qwen3_5", 16, None), ("compact", "float32", 2, "qwen3_5", 16, None), ("compact", "float32", 15, "qwen3_5", 16, None), @@ -26,7 +27,7 @@ ("replay", "bfloat16", 2, "qwen3_5", 16, None), ("replay", "bfloat16", 2, "qwen3_5", 4, None), ("replay", "bfloat16", 0, "qwen3_5", 16, None), - ("replay", "float32", 16, "qwen3_5", 16, "capacity"), + ("replay", "float32", 4, "qwen3_5", 4, "capacity"), ("compact", "float32", 2, "llama", 16, "require a GDN"), ("replay", "float32", 2, "llama", 16, "require a GDN"), ], @@ -52,6 +53,34 @@ def test_ssm_mode_validation(monkeypatch, mode, dtype, step, model, capacity, er _validate_ssm_state_mode(args) +@pytest.mark.parametrize("model", ["qwen3_5", "glm5_next"]) +@pytest.mark.parametrize("mode", ["native", "compact", "replay"]) +@pytest.mark.parametrize("step", [15, 16]) +def test_gdn_kda_mtp_width_limit(monkeypatch, model, mode, step): + monkeypatch.setattr("lightllm.utils.config_utils.get_model_type", lambda _: model) + args = StartArgs(ssm_state_mode=mode, mtp_step=step) + if step <= 15: + _validate_ssm_state_mode(args) + else: + with pytest.raises(AssertionError, match="GDN/KDA requires mtp_step <= 15"): + _validate_ssm_state_mode(args) + + +@pytest.mark.parametrize("model", ["qwen3_5", "glm5_next"]) +@pytest.mark.parametrize("compute_mode", ["inline", "precompute", "parallel"]) +@pytest.mark.parametrize("cache_len", [32, 64]) +def test_replay_large_cache_does_not_bypass_mtp_width_limit(monkeypatch, model, compute_mode, cache_len): + monkeypatch.setattr("lightllm.utils.config_utils.get_model_type", lambda _: model) + args = StartArgs( + ssm_state_mode="replay", + ssm_replay_config={"cache_len": cache_len, "compute_mode": compute_mode}, + linear_att_ssm_data_type="bfloat16", + mtp_step=16, + ) + with pytest.raises(AssertionError, match="GDN/KDA requires mtp_step <= 15"): + _validate_ssm_state_mode(args) + + def test_ssm_mode_defaults_and_removed_flag(): parser = make_argument_parser() assert parser.parse_args([]).ssm_state_mode == StartArgs().ssm_state_mode == "native" From 0fff948cbc2be587176a961ed80fbec8f3363d70 Mon Sep 17 00:00:00 2001 From: sufubao Date: Sat, 17 Oct 2026 16:18:09 +0800 Subject: [PATCH 19/40] refactor: centralize recurrent state lifecycle and checkpoints --- .../common/basemodel/attention/linear/gdn.py | 38 ++- .../common/basemodel/attention/linear/kda.py | 36 +-- lightllm/common/basemodel/basemodel.py | 2 +- lightllm/common/basemodel/batch_objs.py | 9 +- lightllm/common/basemodel/infer_struct.py | 7 +- .../qwen3next_mem_manager.py | 44 +--- lightllm/common/recurrent_state.py | 223 ++++++++++++++++++ lightllm/common/req_manager/hybrid_base.py | 2 +- lightllm/common/req_manager/linear_att.py | 171 ++------------ .../mode_backend/chunked_prefill/impl.py | 2 +- .../mode_backend/dp_backend/impl.py | 8 +- .../model_infer/mtp_speculative/utils.py | 4 +- .../basemodel/attention/linear/test_gdn.py | 8 +- .../linear_att/test_acceptance_fusion.py | 51 ++-- .../linear_att/test_replayssm.py | 55 +++-- .../linear_att/test_replayssm_execution.py | 14 +- .../linear_att/test_ssm_autotune.py | 14 +- unit_tests/common/test_recurrent_state.py | 112 +++++++++ unit_tests/models/glm5_next/test_cache.py | 28 ++- unit_tests/models/glm5_next/test_mtp.py | 26 +- unit_tests/models/glm5_next/test_pd_cache.py | 52 ++-- unit_tests/models/test_post_layer_output.py | 2 +- 22 files changed, 552 insertions(+), 356 deletions(-) create mode 100644 lightllm/common/recurrent_state.py create mode 100644 unit_tests/common/test_recurrent_state.py diff --git a/lightllm/common/basemodel/attention/linear/gdn.py b/lightllm/common/basemodel/attention/linear/gdn.py index 5f1c988e45..571f85d8ac 100644 --- a/lightllm/common/basemodel/attention/linear/gdn.py +++ b/lightllm/common/basemodel/attention/linear/gdn.py @@ -120,9 +120,7 @@ def init_state(self): # 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 * self.backend.req_manager.ssm_slots_per_req - if self.backend.req_manager.ssm_update_cache is not None: - self.backend.req_manager.ssm_update_cache.merge_accepted_updates(self.infer_state.b_req_idx) + self.b_ssm_buffer_idx = self.backend.req_manager.recurrent_state.prepare_prefill(self.infer_state.b_req_idx) return def prefill_att( @@ -217,11 +215,9 @@ def init_state(self): self._init_dynamic_mtp_decode_state(draft_step + 1) else: self._init_fixed_mtp_decode_state(draft_step) - ssm_updates = self.backend.req_manager.ssm_update_cache - if ssm_updates is not None: - ssm_updates.prepare_decode(self.b_conv_buffer_idx, self.b1_mtp_cu_q_seq_len) - if draft_step == 0: - ssm_updates.accept_updates(self.b_conv_buffer_idx) + self.backend.req_manager.recurrent_state.prepare_decode( + self.b_conv_buffer_idx, self.b1_mtp_cu_q_seq_len, single_token=draft_step == 0 + ) def _init_normal_decode_state(self): self.b_conv_buffer_idx = self.infer_state.b_req_idx @@ -235,7 +231,7 @@ def _init_dynamic_mtp_decode_state(self, mtp_size: int): ) = build_dynamic_mtp_linear_att_state_params( b_req_idx=self.infer_state.b_req_idx, b_mtp_index=self.infer_state.b_mtp_index, - req_to_mtp_state_index=self.backend.req_manager.req_to_mtp_state_index, + req_to_mtp_state_index=self.backend.req_manager.recurrent_state.accepted_index, hold_req_id=self.backend.req_manager.HOLD_REQUEST_ID, ) self._init_mtp_ssm_buffer_idx(mtp_size) @@ -257,11 +253,11 @@ def _init_fixed_mtp_decode_state(self, draft_step: int): device=self.infer_state.b_req_idx.device, ) self.b_conv_buffer_idx = self.infer_state.b_req_idx.view(att_batch_size, mtp_size)[:, 0].contiguous() - self.b_num_accepted_tokens = self.backend.req_manager.req_to_mtp_state_index[self.b_conv_buffer_idx] + 1 + self.b_num_accepted_tokens = self.backend.req_manager.recurrent_state.accepted_index[self.b_conv_buffer_idx] + 1 self._init_mtp_ssm_buffer_idx(mtp_size) def _init_mtp_ssm_buffer_idx(self, mtp_size: int): - if self.backend.req_manager.ssm_update_cache is not None: + if self.backend.req_manager.recurrent_state.updates is not None: self.b_ssm_buffer_idx = self.b_conv_buffer_idx return att_batch_size = self.b_conv_buffer_idx.shape[0] @@ -351,17 +347,19 @@ def _gdn_decode_kernel( backend.tp_num_v_heads, backend.head_v_dim, ) - ssm_updates = backend.req_manager.ssm_update_cache - if ssm_updates is not None: + recurrent_state = backend.req_manager.recurrent_state + if recurrent_state.updates is not None: layer = ( layer_weight.layer_num_ - layer_weight.layer_num_ // backend.req_manager.linear_config.full_attention_interval ) args = (query, key, value, a, b, layer_weight.linear_A_log.weight, layer_weight.linear_dt_bias.weight) if layer == 0: - self.infer_state.ssm_run_config = ssm_updates.get_run_config(*args) + self.infer_state.recurrent_state_context = recurrent_state.decode_context(*args) return ( - ssm_updates.forward(layer, *args, self.b_conv_buffer_idx, run_config=self.infer_state.ssm_run_config), + recurrent_state.forward_updates( + layer, *args, self.b_conv_buffer_idx, context=self.infer_state.recurrent_state_context + ), z, ) core_attn_out, _ = fused_recurrent_gated_delta_rule( @@ -409,21 +407,21 @@ def _gdn_mtp_kernel( ) query, key, value = backend._rearrange_mixed_qkv(mixed_qkv, decode=False) - ssm_updates = backend.req_manager.ssm_update_cache - if ssm_updates is not None: + recurrent_state = backend.req_manager.recurrent_state + if recurrent_state.updates is not None: layer = ( layer_weight.layer_num_ - layer_weight.layer_num_ // backend.req_manager.linear_config.full_attention_interval ) args = (query, key, value, a, b, layer_weight.linear_A_log.weight, layer_weight.linear_dt_bias.weight) if layer == 0: - self.infer_state.ssm_run_config = ssm_updates.get_run_config(*args, cu_seqlens_q) - return ssm_updates.forward( + self.infer_state.recurrent_state_context = recurrent_state.decode_context(*args, cu_seqlens_q) + return recurrent_state.forward_updates( layer, *args, self.b_conv_buffer_idx, cu_seqlens_q, - run_config=self.infer_state.ssm_run_config, + context=self.infer_state.recurrent_state_context, ) 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, diff --git a/lightllm/common/basemodel/attention/linear/kda.py b/lightllm/common/basemodel/attention/linear/kda.py index 6b36311879..a896219d6a 100644 --- a/lightllm/common/basemodel/attention/linear/kda.py +++ b/lightllm/common/basemodel/attention/linear/kda.py @@ -45,8 +45,8 @@ def __init__(self, model: "TpPartBaseModel"): self.conv_kernel_size = config["short_conv_kernel_size"] self.lower_bound = config.get("gate_lower_bound", -5.0) self.mtp_step = model.args.mtp_step - if model.req_manager.ssm_update_cache is not None: - model.req_manager.ssm_update_cache.lower_bound = self.lower_bound + if model.req_manager.recurrent_state.updates is not None: + model.req_manager.recurrent_state.updates.lower_bound = self.lower_bound def create_att_prefill_state(self, infer_state: "InferStateInfo"): return KDAPrefillAttState(backend=self, infer_state=infer_state) @@ -66,9 +66,7 @@ class KDAPrefillAttState(BasePrefillAttState): def init_state(self): self.b_conv_buffer_idx = self.infer_state.b_req_idx manager = self.infer_state.req_manager - self.b_ssm_buffer_idx = self.infer_state.b_req_idx * manager.ssm_slots_per_req - if manager.ssm_update_cache is not None: - manager.ssm_update_cache.merge_accepted_updates(self.infer_state.b_req_idx) + self.b_ssm_buffer_idx = manager.recurrent_state.prepare_prefill(self.infer_state.b_req_idx) def prefill_att( self, @@ -142,11 +140,9 @@ def init_state(self): self._init_dynamic_mtp_decode_state(mtp_step + 1) else: self._init_fixed_mtp_decode_state(mtp_step) - updates = self.infer_state.req_manager.ssm_update_cache - if updates is not None: - updates.prepare_decode(self.b_conv_buffer_idx, self.b1_mtp_cu_q_seq_len) - if mtp_step == 0: - updates.accept_updates(self.b_conv_buffer_idx) + self.infer_state.req_manager.recurrent_state.prepare_decode( + self.b_conv_buffer_idx, self.b1_mtp_cu_q_seq_len, single_token=mtp_step == 0 + ) def _init_normal_decode_state(self): self.b_conv_buffer_idx = self.infer_state.b_req_idx @@ -160,7 +156,7 @@ def _init_dynamic_mtp_decode_state(self, mtp_size: int): ) = build_dynamic_mtp_linear_att_state_params( b_req_idx=self.infer_state.b_req_idx, b_mtp_index=self.infer_state.b_mtp_index, - req_to_mtp_state_index=self.infer_state.req_manager.req_to_mtp_state_index, + req_to_mtp_state_index=self.infer_state.req_manager.recurrent_state.accepted_index, hold_req_id=self.infer_state.req_manager.HOLD_REQUEST_ID, ) self._init_mtp_ssm_buffer_idx(mtp_size) @@ -182,11 +178,13 @@ def _init_fixed_mtp_decode_state(self, mtp_step: int): device=self.infer_state.b_req_idx.device, ) self.b_conv_buffer_idx = self.infer_state.b_req_idx.view(att_batch_size, mtp_size)[:, 0].contiguous() - self.b_num_accepted_tokens = self.infer_state.req_manager.req_to_mtp_state_index[self.b_conv_buffer_idx] + 1 + self.b_num_accepted_tokens = ( + self.infer_state.req_manager.recurrent_state.accepted_index[self.b_conv_buffer_idx] + 1 + ) self._init_mtp_ssm_buffer_idx(mtp_size) def _init_mtp_ssm_buffer_idx(self, mtp_size: int): - if self.infer_state.req_manager.ssm_update_cache is not None: + if self.infer_state.req_manager.recurrent_state.updates is not None: self.b_ssm_buffer_idx = self.b_conv_buffer_idx return att_batch_size = self.b_conv_buffer_idx.shape[0] @@ -234,8 +232,8 @@ def decode_att( v = v.view(*shape, backend.tp_num_heads, backend.head_dim) raw_gate = raw_gate.view(*shape, backend.tp_hidden_size) raw_beta = raw_beta.view(*shape, backend.tp_num_heads) - updates = self.infer_state.req_manager.ssm_update_cache - if updates is not None: + recurrent_state = self.infer_state.req_manager.recurrent_state + if recurrent_state.updates is not None: layer = layer_num - layer_num // self.infer_state.req_manager.linear_config.full_attention_interval args = ( q, @@ -247,13 +245,15 @@ def decode_att( layer_weight.linear_dt_bias.weight, ) if layer == 0: - self.infer_state.ssm_run_config = updates.get_run_config(*args, self.b1_mtp_cu_q_seq_len) - return updates.forward( + self.infer_state.recurrent_state_context = recurrent_state.decode_context( + *args, self.b1_mtp_cu_q_seq_len + ) + return recurrent_state.forward_updates( layer, *args, self.b_conv_buffer_idx, self.b1_mtp_cu_q_seq_len, - run_config=self.infer_state.ssm_run_config, + context=self.infer_state.recurrent_state_context, ) output, _ = fused_recurrent_kda( q=q, diff --git a/lightllm/common/basemodel/basemodel.py b/lightllm/common/basemodel/basemodel.py index d328fbd501..f9866b13fe 100755 --- a/lightllm/common/basemodel/basemodel.py +++ b/lightllm/common/basemodel/basemodel.py @@ -530,7 +530,7 @@ def _create_model_output(self, post_output: PostLayerOutput, infer_state: InferS mtp_collector=infer_state.hidden_collector.finish_output(infer_state=infer_state), prompt_logics=infer_state.prompt_logics, logits_token_ids=post_output.logits_token_ids, - ssm_run_config=infer_state.ssm_run_config, + recurrent_state_context=infer_state.recurrent_state_context, logits_row_stride=infer_state.logits_row_stride, ) return output diff --git a/lightllm/common/basemodel/batch_objs.py b/lightllm/common/basemodel/batch_objs.py index f73c2dc153..52218c143a 100644 --- a/lightllm/common/basemodel/batch_objs.py +++ b/lightllm/common/basemodel/batch_objs.py @@ -1,11 +1,14 @@ import copy from dataclasses import dataclass -from typing import List, Optional +from typing import TYPE_CHECKING, List, Optional import torch from lightllm.utils.tensor_utils import tensor_to_no_ref_tensor +if TYPE_CHECKING: + from lightllm.common.recurrent_state import RecurrentStateContext + @dataclass class ModelInput: @@ -207,8 +210,8 @@ class ModelOutput: # target 模型会将候选值回填到完整词表 logits，因此不需要该映射。 logits_token_ids: Optional[torch.Tensor] = None - # Layout selected by this verify graph; preserved when logits are unpadded. - ssm_run_config: Optional[dict] = None + # Opaque commit context for this forward/graph; preserved when unpadding. + recurrent_state_context: Optional["RecurrentStateContext"] = None # Number of physical decode rows represented by each selected logits row. logits_row_stride: int = 1 diff --git a/lightllm/common/basemodel/infer_struct.py b/lightllm/common/basemodel/infer_struct.py index a88c629a81..4fada1e7d8 100755 --- a/lightllm/common/basemodel/infer_struct.py +++ b/lightllm/common/basemodel/infer_struct.py @@ -4,7 +4,7 @@ from lightllm.common.kv_cache_mem_manager import MemoryManager from lightllm.common.req_manager import ReqManager from lightllm.distributed import CustomProcessGroup -from typing import Tuple, Any, Optional, List +from typing import TYPE_CHECKING, Tuple, Any, Optional, List from .triton_kernel.gen_prefill_params import gen_prefill_params from .triton_kernel.gen_decode_params import gen_decode_params from .triton_kernel.multimodal_emb import mark_multimodal_obj @@ -13,6 +13,9 @@ from lightllm.utils.dist_utils import get_global_dp_rank, get_dp_world_size from .attention import BasePrefillAttState, BaseDecodeAttState +if TYPE_CHECKING: + from lightllm.common.recurrent_state import RecurrentStateContext + class InferStateInfo: """ @@ -64,7 +67,7 @@ def __init__(self): # 在microbatch overlap的运行模式下，用于标记当前 microbatch 的 index 序号 # 在一些细节场景下需要有该信息区分一些资源的申请和管理。 self.microbatch_index: int = 0 - self.ssm_run_config: Optional[dict] = None + self.recurrent_state_context: Optional["RecurrentStateContext"] = None self.b_logits_rows: Optional[torch.Tensor] = None self.logits_row_stride: int = 1 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 0e4e3364c9..f9dae5b25f 100644 --- a/lightllm/common/kv_cache_mem_manager/qwen3next_mem_manager.py +++ b/lightllm/common/kv_cache_mem_manager/qwen3next_mem_manager.py @@ -71,11 +71,7 @@ def _free_linear_att_buffers(self): return def write_to_shm(self, req_manager): - self.ssm_update_cache = req_manager.ssm_update_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 + self.recurrent_state = req_manager.recurrent_state # super().write_to_shm() 会用 ForkingPickler 序列化本对象，torch 在 dump 时会把 # CPU tensor 的 storage 原地迁到共享内存，使本进程大页 state cache 原本 # pinned(cudaHostAlloc) 的内存退化为普通 shm mmap，之后 Triton kernel 携带该指针 @@ -176,8 +172,6 @@ class Qwen3NextLinearAttPageHelper: def __init__(self, mem_manager: "Qwen3NextMemManager"): self.mem_manager = mem_manager self.linear_config = mem_manager.linear_config - self.req_to_conv_state = mem_manager.req_to_conv_state - self.req_to_ssm_state = mem_manager.req_to_ssm_state self.global_linear_k_heads = self.linear_config.global_linear_k_heads self.global_linear_v_heads = self.linear_config.global_linear_v_heads @@ -199,9 +193,9 @@ def __init__(self, mem_manager: "Qwen3NextMemManager"): ) self.conv_nbytes = ( - self.conv_shape[0] * self.conv_shape[1] * self.conv_shape[2] * self.req_to_conv_state.buffer.element_size() + self.conv_shape[0] * self.conv_shape[1] * self.conv_shape[2] * self.linear_config.conv_state_dtype.itemsize ) - ssm_alignment = self.req_to_ssm_state.buffer.element_size() + ssm_alignment = self.linear_config.ssm_state_dtype.itemsize # 做一下字节对齐，防止切出来的不对齐，导致一些操作的性能下降。 self.ssm_offset = ((self.conv_nbytes + ssm_alignment - 1) // ssm_alignment) * ssm_alignment self.ssm_nbytes = ( @@ -209,7 +203,7 @@ def __init__(self, mem_manager: "Qwen3NextMemManager"): * self.ssm_shape[1] * self.ssm_shape[2] * self.ssm_shape[3] - * self.req_to_ssm_state.buffer.element_size() + * self.linear_config.ssm_state_dtype.itemsize ) self.state_nbytes = self.ssm_offset + self.ssm_nbytes @@ -225,7 +219,7 @@ def get_dp_mems(self, mem_managers: List["Qwen3NextMemManager"], dp_index: int, dp_mems = mem_managers[(dp_index * dp_world_size) : ((dp_index + 1) * dp_world_size)] assert len(dp_mems) == dp_world_size for mem in dp_mems: - assert hasattr(mem, "req_to_conv_state") and hasattr(mem, "req_to_ssm_state") + assert hasattr(mem, "recurrent_state") assert mem.linear_config.linear_layer_num == self.linear_config.linear_layer_num assert mem.linear_config.conv_kernel_size == self.linear_config.conv_kernel_size assert mem.linear_config.head_linear_k_dim == self.linear_config.head_linear_k_dim @@ -236,10 +230,10 @@ def get_dp_mems(self, mem_managers: List["Qwen3NextMemManager"], dp_index: int, def view_page_to_linear_att_state(self, page_index: int): page_bytes = self.mem_manager.kv_move_buffer[page_index].view(torch.uint8).reshape(-1) - conv_page = page_bytes[0 : self.conv_nbytes].view(self.req_to_conv_state.buffer.dtype).view(self.conv_shape) + conv_page = page_bytes[0 : self.conv_nbytes].view(self.linear_config.conv_state_dtype).view(self.conv_shape) ssm_page = ( page_bytes[self.ssm_offset : self.ssm_offset + self.ssm_nbytes] - .view(self.req_to_ssm_state.buffer.dtype) + .view(self.linear_config.ssm_state_dtype) .view(self.ssm_shape) ) return conv_page, ssm_page @@ -266,9 +260,6 @@ def read_page_to_req( self._read_one_rank(mem, tp_index, req_idx, conv_page, ssm_page) return - def _get_req_state_indexes(self, req_idx: int): - return req_idx, req_idx * self.mem_manager.ssm_slots_per_req - def _write_one_rank( self, mem: "Qwen3NextMemManager", @@ -277,17 +268,7 @@ def _write_one_rank( conv_page: torch.Tensor, ssm_page: torch.Tensor, ): - ssm_updates = mem.ssm_update_cache - 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 ssm_updates 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 = ( - ssm_updates.snapshot_accepted_state(req_idx) - if ssm_updates is not None - else mem.req_to_ssm_state.buffer[:, ssm_req_idx, ...] - ) + conv_state, ssm_state = mem.recurrent_state.snapshot(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) return @@ -465,14 +446,7 @@ def _read_one_rank( conv_page: torch.Tensor, ssm_page: torch.Tensor, ): - ssm_updates = mem.ssm_update_cache - if ssm_updates is not None: - ssm_updates.clear_history(req_idx) - if mem.req_to_mtp_state_index is not None: - mem.req_to_mtp_state_index[req_idx].zero_() - 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, ...] + conv_state, ssm_state = mem.recurrent_state.restore_target(req_idx) self._copy_page_to_conv_state(conv_page, conv_state, mem, tp_index) self._copy_page_to_ssm_state(ssm_page, ssm_state, mem, tp_index) return diff --git a/lightllm/common/recurrent_state.py b/lightllm/common/recurrent_state.py new file mode 100644 index 0000000000..e0bb208efa --- /dev/null +++ b/lightllm/common/recurrent_state.py @@ -0,0 +1,223 @@ +"""Worker-local recurrent state shared by attention, checkpoints, and PD I/O. + +The runner orders operations on its current CUDA stream. This component does +not allocate request IDs, choose prefix boundaries, or synchronize streams. +Attention backends retain responsibility for recurrent computation. +""" + +from dataclasses import dataclass +from typing import NamedTuple, Optional + +import torch + +from lightllm.common.state_cache_manager import LayerCache, LinearAttCacheConfig + + +@dataclass(frozen=True) +class RecurrentStateContext: + """Commit information for one forward/graph, passed unchanged by the runner. + + Compact reconstruction must use the verify kernel's reduction layout. + Keeping it per forward also permits overlapping graphs with different layouts. + """ + + _run_config: Optional[dict] = None + + +class RecurrentStateSnapshot(NamedTuple): + conv: torch.Tensor + ssm: torch.Tensor + + +class RecurrentState: + def __init__(self, max_request_num, linear_config: LinearAttCacheConfig, args, *, kda=False, device="cuda"): + self.linear_config = linear_config + self.mtp_step = args.mtp_step + self.slots_per_req = self.mtp_step + 1 if args.ssm_state_mode == "native" else 1 + self.accepted_index = ( + torch.zeros(max_request_num + 1, dtype=torch.int32, device=device) if self.mtp_step > 0 else None + ) + self.conv_state = LayerCache( + size=max_request_num + 1, + dtype=linear_config.conv_state_dtype, + shape=linear_config.get_mtp_conv_state_shape(mtp_step=self.mtp_step), + layer_num=linear_config.linear_layer_num, + device=device, + ) + self.ssm_state = LayerCache( + size=(max_request_num + 1) * self.slots_per_req, + dtype=linear_config.ssm_state_dtype, + shape=linear_config.get_ssm_state_shape(), + layer_num=linear_config.linear_layer_num, + device=device, + ) + self.updates = None + if args.ssm_state_mode == "replay": + from lightllm.common.basemodel.triton_kernel.linear_att.replayssm import ReplaySSMCache + + self.updates = ReplaySSMCache( + self.ssm_state.buffer, + args.ssm_replay_config.get("cache_len", 16), + self.mtp_step + 1, + linear_config.conv_state_dtype, + num_key_heads=linear_config.num_linear_k_heads, + kda=kda, + projection_mode=args.ssm_replay_config.get("compute_mode", "inline"), + fold_method=args.ssm_replay_config.get("fold_method", "raw"), + ) + elif args.ssm_state_mode == "compact": + from lightllm.common.basemodel.triton_kernel.linear_att.replayssm_compact import CompactSSMCache + + self.updates = CompactSSMCache( + self.ssm_state.buffer, + self.mtp_step + 1, + linear_config.conv_state_dtype, + num_key_heads=linear_config.num_linear_k_heads, + kda=kda, + ) + + def _clear_positions(self, req_idx): + if self.updates is not None: + self.updates.clear_history(req_idx) + if self.accepted_index is not None: + self.accepted_index[req_idx].zero_() + + def reset(self, req_idx): + """Initialize a newly allocated/reused request, including speculative slots.""" + self._clear_positions(req_idx) + self.conv_state.buffer[:, req_idx].zero_() + start = req_idx * self.slots_per_req + self.ssm_state.buffer[:, start : start + self.slots_per_req].zero_() + + def prepare_prefill(self, reqs): + if self.updates is not None: + self.updates.merge_accepted_updates(reqs) + return reqs * self.slots_per_req + + def prepare_decode(self, reqs, cu_seqlens=None, *, single_token=False): + if self.updates is not None: + self.updates.prepare_decode(reqs, cu_seqlens) + if single_token: + # Preserve the forward cursor snapshot before accepting the + # guaranteed token, including when forward is graph-replayed. + self.updates.accept_updates(reqs) + + def decode_context(self, *args): + return RecurrentStateContext(self.updates.get_run_config(*args)) + + def forward_updates(self, layer, *args, context): + return self.updates.forward(layer, *args, run_config=context._run_config) + + def commit(self, b_req_mtp_start_loc, b_req_idx, b_mtp_index, accepted_index, verify_width, batches=None): + """Commit accepted prefixes without requiring full-state materialization. + + ``batches`` partitions logical requests in the same order as the MTP + start locations. Each context belongs to the forward that wrote that + partition's scratch; padding rows do not participate in acceptance. + """ + from lightllm.common.basemodel.triton_kernel.mtp_utils import linear_att_mtp_state_index_update + from lightllm.common.basemodel.triton_kernel.linear_att.replayssm import ReplaySSMCache + + replay_cursors = None + if isinstance(self.updates, ReplaySSMCache) and self.updates.verify_width > 1: + replay_cursors = self.updates.cursors + linear_att_mtp_state_index_update( + req_to_mtp_state_index=self.accepted_index, + b_req_mtp_start_loc=b_req_mtp_start_loc, + b_req_idx=b_req_idx, + b_mtp_index=b_mtp_index, + accepted_index=accepted_index, + verify_width=verify_width, + replay_cursors=replay_cursors, + ) + if self.updates is None or replay_cursors is not None: + return + reqs = b_req_idx[b_req_mtp_start_loc.long()] + if batches is None or isinstance(self.updates, ReplaySSMCache): + self.updates.accept_updates(reqs, self.accepted_index) + return + assert sum(size for size, _ in batches) == reqs.numel() + start = 0 + for size, context in batches: + if size: + self.updates.accept_updates( + reqs[start : start + size], + self.accepted_index, + run_config=context._run_config if context is not None else None, + ) + start += size + + def snapshot(self, req_idx): + """Export the current accepted prefix without mutating active state. + + The caller associates this prefix with its cache/transport token boundary. + These tensors may alias the active pool: enqueue consuming copies on the + current stream before advancing/releasing the request. Cross-stream or + cross-process callers must use their existing completion events before + recycling sources or destinations. No CPU transfer or sync is implicit. + """ + width = self.linear_config.get_conv_state_shape()[-1] + conv = self.conv_state.buffer[:, req_idx, ..., :width] + if self.accepted_index is not None: + offsets = torch.arange(width, device=conv.device) + self.accepted_index[req_idx] + conv = self.conv_state.buffer[:, req_idx].index_select(-1, offsets) + if self.updates is not None: + ssm = self.updates.snapshot_accepted_state(req_idx) + elif self.slots_per_req > 1: + index = req_idx * self.slots_per_req + self.accepted_index[req_idx] + ssm = self.ssm_state.buffer.index_select(1, index.long().view(1))[:, 0] + else: + ssm = self.ssm_state.buffer[:, req_idx] + return RecurrentStateSnapshot(conv, ssm) + + def restore_target(self, req_idx): + """Invalidate old history and return canonical destinations for restore. + + The caller must fill both tensors before the next forward. This lets PD + transport unpack TP shards directly into the pool without a staging copy. + The same stream/completion rules as ``snapshot`` apply. + """ + self._clear_positions(req_idx) + width = self.linear_config.get_conv_state_shape()[-1] + return RecurrentStateSnapshot( + self.conv_state.buffer[:, req_idx, ..., :width], + self.ssm_state.buffer[:, req_idx * self.slots_per_req], + ) + + def restore(self, req_idx, conv, ssm): + target = self.restore_target(req_idx) + target.conv.copy_(conv, non_blocking=True) + target.ssm.copy_(ssm, non_blocking=True) + + def save_checkpoints(self, b_req_idx, req_indexes, buffer_indexes, destination): + """Copy current accepted prefixes into caller-allocated checkpoint slots.""" + if self.updates is None: + from lightllm.common.basemodel.triton_kernel.linear_att_copy import copy_linear_att_state_to_kv_buffer + + indexes = torch.tensor(buffer_indexes, dtype=torch.int32, device="cpu").to( + b_req_idx.device, non_blocking=True + ) + copy_linear_att_state_to_kv_buffer( + b_req_idx=b_req_idx, + big_page_buffer_ids=indexes, + gpu_conv_state=self.conv_state.buffer, + gpu_ssm_state=self.ssm_state.buffer, + cpu_kv_conv_state=destination.conv_state_cache.buffer, + cpu_kv_ssm_state=destination.ssm_state_cache.buffer, + mtp_step=self.slots_per_req - 1, + conv_offsets=self.accepted_index, + ) + else: + for req_idx, buffer_idx in zip(req_indexes, buffer_indexes): + if buffer_idx != -1: + self.export(req_idx, *destination.get_state_cache(buffer_idx)) + + def export(self, req_idx, conv, ssm): + snapshot = self.snapshot(req_idx) + conv.copy_(snapshot.conv, non_blocking=True) + ssm.copy_(snapshot.ssm, non_blocking=True) + + def layer_buffers(self, layer_idx_in_all): + assert 0 <= layer_idx_in_all < self.linear_config.all_layer_num + layer = layer_idx_in_all - layer_idx_in_all // self.linear_config.full_attention_interval + return self.conv_state.buffer[layer], self.ssm_state.buffer[layer] diff --git a/lightllm/common/req_manager/hybrid_base.py b/lightllm/common/req_manager/hybrid_base.py index e71e86edb5..370c8c4161 100644 --- a/lightllm/common/req_manager/hybrid_base.py +++ b/lightllm/common/req_manager/hybrid_base.py @@ -66,7 +66,7 @@ def save_state(self, req_idx: int, buffer_idx: int, state_cache_manager): """请求 GPU 运行态 → 指定 CPU checkpoint 槽位；大小页共用，调用方负责分配槽位。""" def update_mtp_state( - self, b_req_mtp_start_loc, b_req_idx, b_mtp_index, accepted_index, verify_width, ssm_accept_batches=None + self, b_req_mtp_start_loc, b_req_idx, b_mtp_index, accepted_index, verify_width, recurrent_state_batches=None ): """接受推测 token 后更新运行态位置；需要调整状态索引的模型覆写。""" return diff --git a/lightllm/common/req_manager/linear_att.py b/lightllm/common/req_manager/linear_att.py index 6a96e2ac22..9499ab2ee4 100644 --- a/lightllm/common/req_manager/linear_att.py +++ b/lightllm/common/req_manager/linear_att.py @@ -2,7 +2,8 @@ import torch -from lightllm.common.state_cache_manager import LayerCache, LinearAttCacheConfig, LinearAttCacheManager +from lightllm.common.recurrent_state import RecurrentState +from lightllm.common.state_cache_manager import LinearAttCacheConfig, LinearAttCacheManager from lightllm.utils.envs_utils import get_env_start_args from .hybrid_base import HybridAttentionReqManager @@ -15,180 +16,40 @@ class ReqManagerForMamba(HybridAttentionReqManager): kda = False - 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): super().__init__(max_request_num, max_sequence_length, mem_manager) args = get_env_start_args() self.mtp_step = args.mtp_step - self.ssm_slots_per_req = self.mtp_step + 1 if args.ssm_state_mode == "native" else 1 - self.ssm_update_cache = None - # 因为在mtp的推理中，需要标记每个请求对应的mtp index状态(conv state 和 ssm state)，在mtp对应序列中 - # 的真实位置，所以需要需要一个标记来记录，不然算子无法找到真实的处理起点。 - self.req_to_mtp_state_index = ( - torch.zeros((max_request_num + 1,), dtype=torch.int32, device="cuda") if self.mtp_step > 0 else None - ) - # 突然想到， 在linear att 开启mtp的模式中，现在的prefill linear att 算子默认是从0的位置读取信息进行操作 - # 所以不能支持 prefill decode mixed 操作了，因为一个decode过的请求，重新用prefill 算子跑，会出现读错linear - # 状态位置的问题。导致bug, 在这里加个断言，以后可以支持上 TODO + # Prefill kernels read the canonical convolution window. Mixed prefill + # and MTP decode still require a separate execution-layout change. if self.mtp_step > 0: - assert get_env_start_args().enable_prefill_decode_mixed is False - - self.big_page_token_num = ( - get_env_start_args().linear_att_page_block_num * get_env_start_args().linear_att_hash_page_size - ) + assert args.enable_prefill_decode_mixed is False + self.big_page_token_num = args.linear_att_page_block_num * args.linear_att_hash_page_size self.linear_config = linear_config - - self.req_to_conv_state = LayerCache( - size=(max_request_num + 1), - dtype=self.linear_config.conv_state_dtype, - shape=self.linear_config.get_mtp_conv_state_shape(mtp_step=self.mtp_step), - layer_num=self.linear_config.linear_layer_num, - device="cuda", - ) - self.req_to_ssm_state = LayerCache( - 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 args.ssm_state_mode == "replay": - from lightllm.common.basemodel.triton_kernel.linear_att.replayssm import ReplaySSMCache - - self.ssm_update_cache = ReplaySSMCache( - self.req_to_ssm_state.buffer, - args.ssm_replay_config.get("cache_len", 16), - self.mtp_step + 1, - linear_config.conv_state_dtype, - num_key_heads=linear_config.num_linear_k_heads, - kda=self.kda, - projection_mode=args.ssm_replay_config.get("compute_mode", "inline"), - fold_method=args.ssm_replay_config.get("fold_method", "raw"), - ) - elif args.ssm_state_mode == "compact": - from lightllm.common.basemodel.triton_kernel.linear_att.replayssm_compact import CompactSSMCache - - self.ssm_update_cache = CompactSSMCache( - self.req_to_ssm_state.buffer, - self.mtp_step + 1, - linear_config.conv_state_dtype, - num_key_heads=linear_config.num_linear_k_heads, - kda=self.kda, - ) - return + self.recurrent_state = RecurrentState(max_request_num, linear_config, args, kda=self.kda) def init_hybrid_attention_state(self, req: "InferReq"): - if self.ssm_update_cache is not None: - # 清除这个请求槽位之前的历史记录标记 - self.ssm_update_cache.clear_history(req.req_idx) - conv_index = req.req_idx - ssm_start = req.req_idx * self.ssm_slots_per_req - self.req_to_conv_state.buffer[:, conv_index, ...].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].zero_() - return + self.recurrent_state.reset(req.req_idx) def create_small_page_cache_manager(self, size: int): self.small_page_buffers = LinearAttCacheManager(size=size, linear_config=self.linear_config) return self.small_page_buffers 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.ssm_update_cache is not None: - return super().save_big_page_states(b_req_idx, req_indexes, buffer_indexes) - 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( - b_req_idx=b_req_idx, - big_page_buffer_ids=buffer_indexes, - gpu_conv_state=self.req_to_conv_state.buffer, - 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.ssm_slots_per_req - 1, - conv_offsets=self.req_to_mtp_state_index, - ) - return + self.recurrent_state.save_checkpoints(b_req_idx, req_indexes, buffer_indexes, self.big_page_buffers) def save_state(self, req_idx: int, buffer_idx: int, state_cache_manager: LinearAttCacheManager): - # 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] - if 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) - if self.ssm_update_cache is not None: - gpu_ssm_state = self.ssm_update_cache.snapshot_accepted_state(req_idx) - else: - gpu_ssm_state = self.req_to_ssm_state.buffer[:, req_idx * self.ssm_slots_per_req, ...] - if self.ssm_slots_per_req > 1: - state_index = req_idx * self.ssm_slots_per_req + self.req_to_mtp_state_index[req_idx] - gpu_ssm_state = self.req_to_ssm_state.buffer.index_select(1, state_index.long().view(1))[:, 0] - 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) + self.recurrent_state.export(req_idx, *state_cache_manager.get_state_cache(buffer_idx)) def get_mamba_cache(self, layer_idx_in_all: int): - assert ( - 0 <= layer_idx_in_all < self.linear_config.all_layer_num - ), f"invalid transformer layer index {layer_idx_in_all}" - layer_idx_in_linear = layer_idx_in_all - (layer_idx_in_all // self.linear_config.full_attention_interval) - conv_states = self.req_to_conv_state.buffer[layer_idx_in_linear] - ssm_states = self.req_to_ssm_state.buffer[layer_idx_in_linear] - return conv_states, ssm_states + return self.recurrent_state.layer_buffers(layer_idx_in_all) def update_mtp_state( - self, b_req_mtp_start_loc, b_req_idx, b_mtp_index, accepted_index, verify_width, ssm_accept_batches=None + self, b_req_mtp_start_loc, b_req_idx, b_mtp_index, accepted_index, verify_width, recurrent_state_batches=None ): - from lightllm.common.basemodel.triton_kernel.mtp_utils import linear_att_mtp_state_index_update - - replay_cursors = None - if self.ssm_update_cache is not None: - from lightllm.common.basemodel.triton_kernel.linear_att.replayssm import ReplaySSMCache - - if isinstance(self.ssm_update_cache, ReplaySSMCache) and self.ssm_update_cache.verify_width > 1: - replay_cursors = self.ssm_update_cache.cursors - - linear_att_mtp_state_index_update( - req_to_mtp_state_index=self.req_to_mtp_state_index, - b_req_mtp_start_loc=b_req_mtp_start_loc, - b_req_idx=b_req_idx, - b_mtp_index=b_mtp_index, - accepted_index=accepted_index, - verify_width=verify_width, - replay_cursors=replay_cursors, + self.recurrent_state.commit( + b_req_mtp_start_loc, b_req_idx, b_mtp_index, accepted_index, verify_width, recurrent_state_batches ) - if self.ssm_update_cache is not None and replay_cursors is None: - reqs = b_req_idx[b_req_mtp_start_loc.long()] - if ssm_accept_batches is None or isinstance(self.ssm_update_cache, ReplaySSMCache): - self.ssm_update_cache.accept_updates(reqs, self.req_to_mtp_state_index) - else: - assert sum(size for size, _ in ssm_accept_batches) == reqs.numel() - start = 0 - for size, config in ssm_accept_batches: - if size: - self.ssm_update_cache.accept_updates( - reqs[start : start + size], self.req_to_mtp_state_index, run_config=config - ) - start += size - def restore_state(self, req: "InferReq", state_cache_manager: LinearAttCacheManager, buffer_idx: int): - if self.ssm_update_cache is not None: - self.ssm_update_cache.clear_history(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.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 - if self.req_to_mtp_state_index is not None: - self.req_to_mtp_state_index[req.req_idx].zero_() - return + self.recurrent_state.restore(req.req_idx, *state_cache_manager.get_state_cache(buffer_idx)) diff --git a/lightllm/server/router/model_infer/mode_backend/chunked_prefill/impl.py b/lightllm/server/router/model_infer/mode_backend/chunked_prefill/impl.py index 78ae3d57af..30263e5cae 100644 --- a/lightllm/server/router/model_infer/mode_backend/chunked_prefill/impl.py +++ b/lightllm/server/router/model_infer/mode_backend/chunked_prefill/impl.py @@ -295,7 +295,7 @@ def decode_mtp( b_req_idx=model_input.b_req_idx, b_req_mtp_start_loc=b_req_mtp_start_loc, b_mtp_index=model_input.b_mtp_index, - ssm_accept_batches=((req_num, model_output.ssm_run_config),), + recurrent_state_batches=((req_num, model_output.recurrent_state_context),), ) accepted_index_cpu = g_pin_mem_manager.async_copy_from_gpu_tensor( key="accepted_index", diff --git a/lightllm/server/router/model_infer/mode_backend/dp_backend/impl.py b/lightllm/server/router/model_infer/mode_backend/dp_backend/impl.py index 5ac4dec2d3..f4e6e2aee7 100644 --- a/lightllm/server/router/model_infer/mode_backend/dp_backend/impl.py +++ b/lightllm/server/router/model_infer/mode_backend/dp_backend/impl.py @@ -524,7 +524,7 @@ def decode_mtp(self, event_pack: OverlapEventPack, decode_reqs: List[InferReq]): b_req_idx=model_input.b_req_idx, b_req_mtp_start_loc=b_req_mtp_start_loc, b_mtp_index=model_input.b_mtp_index, - ssm_accept_batches=((req_num, model_output.ssm_run_config),), + recurrent_state_batches=((req_num, model_output.recurrent_state_context),), ) accepted_index_cpu = g_pin_mem_manager.async_copy_from_gpu_tensor( key="accepted_index", @@ -799,9 +799,9 @@ def decode_overlap_mtp(self, event_pack: OverlapEventPack, decode_reqs: List[Inf b_req_idx=b_req_idx, b_req_mtp_start_loc=b_req_mtp_start_loc, b_mtp_index=b_mtp_index, - ssm_accept_batches=( - (real_request_num0, model_output0.ssm_run_config), - (real_request_num1, model_output1.ssm_run_config), + recurrent_state_batches=( + (real_request_num0, model_output0.recurrent_state_context), + (real_request_num1, model_output1.recurrent_state_context), ), ) mtp_accept_len0 = mtp_accept_len[:real_request_num0] diff --git a/lightllm/server/router/model_infer/mtp_speculative/utils.py b/lightllm/server/router/model_infer/mtp_speculative/utils.py index 3a3c47c549..df71a9f5e9 100644 --- a/lightllm/server/router/model_infer/mtp_speculative/utils.py +++ b/lightllm/server/router/model_infer/mtp_speculative/utils.py @@ -22,7 +22,7 @@ def verify_mtp_tokens( b_req_idx: torch.Tensor, b_req_mtp_start_loc: torch.Tensor, b_mtp_index: torch.Tensor, - ssm_accept_batches=None, + recurrent_state_batches=None, ) -> Tuple[torch.Tensor, torch.Tensor]: """Verify target tokens and update recurrent MTP state when required.""" @@ -39,7 +39,7 @@ def verify_mtp_tokens( b_mtp_index=b_mtp_index, accepted_index=accepted_index, verify_width=backend.max_draft_step + 1, - ssm_accept_batches=ssm_accept_batches, + recurrent_state_batches=recurrent_state_batches, ) return accept_lengths, accepted_index diff --git a/unit_tests/common/basemodel/attention/linear/test_gdn.py b/unit_tests/common/basemodel/attention/linear/test_gdn.py index de4cfe5040..1b3d8934d1 100644 --- a/unit_tests/common/basemodel/attention/linear/test_gdn.py +++ b/unit_tests/common/basemodel/attention/linear/test_gdn.py @@ -10,6 +10,7 @@ import lightllm.common.basemodel.triton_kernel.linear_att.fla.ops as fla_ops from lightllm.common.basemodel.attention.linear.flashqla import FlashQlaLinearAttBackend from lightllm.common.basemodel.attention.linear.triton import TritonLinearAttBackend +from lightllm.common.recurrent_state import RecurrentState from lightllm.common.req_manager import ReqManager, ReqManagerForMamba from lightllm.server.router.model_infer.infer_batch import InferenceContext from lightllm.common.state_cache_manager import LinearAttCacheConfig @@ -69,9 +70,10 @@ def _create_decode_state(*, draft_step, dynamic_layout, b_req_idx, req_to_mtp_st ) # Only CPU metadata is needed for these layout tests. req_manager = object.__new__(ReqManagerForMamba) + req_manager.recurrent_state = object.__new__(RecurrentState) req_manager.HOLD_REQUEST_ID = req_to_mtp_state_index.shape[0] - 1 if req_to_mtp_state_index is not None else -1 - req_manager.req_to_mtp_state_index = req_to_mtp_state_index - req_manager.ssm_update_cache = None + req_manager.recurrent_state.accepted_index = req_to_mtp_state_index + req_manager.recurrent_state.updates = None infer_state = SimpleNamespace( batch_size=b_req_idx.shape[0], b_req_idx=b_req_idx, @@ -178,7 +180,7 @@ def prepare_decode(actual_reqs, actual_cu): assert actual_reqs is reqs and actual_cu is cu return positions - state.backend.req_manager.ssm_update_cache = SimpleNamespace( + state.backend.req_manager.recurrent_state.updates = SimpleNamespace( prepare_decode=prepare_decode, accept_updates=lambda *args: pytest.fail("speculative updates must wait for acceptance"), ) diff --git a/unit_tests/common/basemodel/triton_kernel/linear_att/test_acceptance_fusion.py b/unit_tests/common/basemodel/triton_kernel/linear_att/test_acceptance_fusion.py index 48de74ec06..83bfa18c5a 100644 --- a/unit_tests/common/basemodel/triton_kernel/linear_att/test_acceptance_fusion.py +++ b/unit_tests/common/basemodel/triton_kernel/linear_att/test_acceptance_fusion.py @@ -3,6 +3,7 @@ import triton from lightllm.common.basemodel.triton_kernel.mtp_utils import linear_att_mtp_state_index_update as update from lightllm.common.basemodel.triton_kernel.linear_att.replayssm import ReplaySSMCache +from lightllm.common.recurrent_state import RecurrentState from lightllm.common.req_manager.linear_att import ReqManagerForMamba @@ -25,14 +26,15 @@ def make(n, width): cache.capacity = max(4, triton.next_power_of_2(width)) cache.verify_width = width manager = ReqManagerForMamba.__new__(ReqManagerForMamba) - manager.req_to_mtp_state_index = tensor([-9] * (hold + 1)) - manager.ssm_update_cache = cache + manager.recurrent_state = object.__new__(RecurrentState) + manager.recurrent_state.accepted_index = tensor([-9] * (hold + 1)) + manager.recurrent_state.updates = cache return manager, args, (ids, lengths, starts) def old(m, args): - update(m.req_to_mtp_state_index, *args) - m.ssm_update_cache.accept_updates(args[1][args[0].long()], m.req_to_mtp_state_index) + update(m.recurrent_state.accepted_index, *args) + m.recurrent_state.updates.accept_updates(args[1][args[0].long()], m.recurrent_state.accepted_index) def new(m, args): @@ -53,19 +55,19 @@ def test_acceptance_fusion_matches_original_and_cpu_under_graph(): graph = torch.cuda.CUDAGraph() with torch.cuda.graph(graph): new(m, args) - expected_index = [-9] * len(m.req_to_mtp_state_index) + expected_index = [-9] * len(m.recurrent_state.accepted_index) expected_cursor = [7] * len(expected_index) - m.req_to_mtp_state_index.fill_(-9) - ref.req_to_mtp_state_index.fill_(-9) - m.ssm_update_cache.cursors.fill_(7) - ref.ssm_update_cache.cursors.fill_(7) + m.recurrent_state.accepted_index.fill_(-9) + ref.recurrent_state.accepted_index.fill_(-9) + m.recurrent_state.updates.cursors.fill_(7) + ref.recurrent_state.updates.cursors.fill_(7) for step in range(24): mask = [] for i, (r, l) in enumerate(zip(ids, lengths)): count = (step + i) % (l + 1) mask.extend([1] * count + [0] * (l - count)) expected_index[r] = count - 1 - if r != m.ssm_update_cache.hold: + if r != m.recurrent_state.updates.hold: expected_cursor[r] += count if width > 1 else 1 args[3].copy_(torch.tensor(mask, device="cuda", dtype=torch.int32)) old(ref, args) @@ -73,15 +75,15 @@ def test_acceptance_fusion_matches_original_and_cpu_under_graph(): graph.replay() else: new(m, args) - assert torch.equal(m.req_to_mtp_state_index, ref.req_to_mtp_state_index) - assert torch.equal(m.ssm_update_cache.cursors, ref.ssm_update_cache.cursors) - assert m.req_to_mtp_state_index.tolist() == expected_index - assert m.ssm_update_cache.cursors.tolist() == expected_cursor + assert torch.equal(m.recurrent_state.accepted_index, ref.recurrent_state.accepted_index) + assert torch.equal(m.recurrent_state.updates.cursors, ref.recurrent_state.updates.cursors) + assert m.recurrent_state.accepted_index.tolist() == expected_index + assert m.recurrent_state.updates.cursors.tolist() == expected_cursor checks += 1 # Native optional-pointer path and manager dispatch leave metadata-only cache absent. - m.ssm_update_cache = None + m.recurrent_state.updates = None new(m, args) - assert m.req_to_mtp_state_index.tolist() == expected_index + assert m.recurrent_state.accepted_index.tolist() == expected_index # Non-Replay caches must still receive the accepted request IDs and indexes. class CompactProbe: @@ -90,10 +92,10 @@ def accept_updates(self, reqs, index): m, args, (ids, _, _) = make(7, 4) probe = CompactProbe() - m.ssm_update_cache = probe + m.recurrent_state.updates = probe new(m, args) assert probe.seen[0].tolist() == ids - assert torch.equal(probe.seen[1], m.req_to_mtp_state_index) + assert torch.equal(probe.seen[1], m.recurrent_state.accepted_index) print("ACCEPT_EQUIVALENCE_CHECKS=" + str(checks), flush=True) @@ -118,8 +120,9 @@ def test_fused_acceptance_preserves_folds_snapshots_and_slot_reuse(dtype, projec managers = [] for cache in caches: manager = ReqManagerForMamba.__new__(ReqManagerForMamba) - manager.ssm_update_cache = cache - manager.req_to_mtp_state_index = torch.zeros(8, dtype=torch.int32, device="cuda") + manager.recurrent_state = object.__new__(RecurrentState) + manager.recurrent_state.updates = cache + manager.recurrent_state.accepted_index = torch.zeros(8, dtype=torch.int32, device="cuda") managers.append(manager) reqs = torch.full((4,), 7, dtype=torch.int32, device="cuda") cu = torch.tensor([0, 3, 6, 9, 12], dtype=torch.int32, device="cuda") @@ -178,7 +181,7 @@ def cycle(i): assert torch.equal(left, right) assert torch.equal(caches[0].state, caches[1].state) assert torch.equal(caches[0].cursors, caches[1].cursors) - assert torch.equal(managers[0].req_to_mtp_state_index, managers[1].req_to_mtp_state_index) + assert torch.equal(managers[0].recurrent_state.accepted_index, managers[1].recurrent_state.accepted_index) if iteration % 7 == 0: assert torch.equal(caches[0].snapshot_accepted_state(3), caches[1].snapshot_accepted_state(3)) if iteration % 11 == 10: @@ -187,9 +190,9 @@ def cycle(i): assert torch.equal(caches[0].state, caches[1].state) if iteration == 19: for manager in managers: - manager.ssm_update_cache.clear_history(1) - manager.ssm_update_cache.state[:, 1].zero_() - manager.req_to_mtp_state_index[1] = 0 + manager.recurrent_state.updates.clear_history(1) + manager.recurrent_state.updates.state[:, 1].zero_() + manager.recurrent_state.accepted_index[1] = 0 for cache in caches: cache.merge_accepted_updates(reqs[:3]) assert torch.equal(caches[0].state, caches[1].state) 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 index 61dfbf791f..bae5eba2cd 100644 --- a/unit_tests/common/basemodel/triton_kernel/linear_att/test_replayssm.py +++ b/unit_tests/common/basemodel/triton_kernel/linear_att/test_replayssm.py @@ -485,6 +485,7 @@ def test_autotune_candidates_match_recurrence(dtype, run_config, mode): @pytest.mark.parametrize("dtype", [torch.float32, torch.bfloat16]) def test_checkpoint_restores_pending_history_and_accepted_conv(save_big, monkeypatch, dtype): from types import SimpleNamespace + from lightllm.common.recurrent_state import RecurrentState 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 @@ -493,6 +494,7 @@ def test_checkpoint_restores_pending_history_and_accepted_conv(save_big, monkeyp state_rtol = 2e-5 if dtype == torch.float32 else torch.finfo(dtype).eps # Real save/restore and PD helpers, with just the unrelated allocator omitted. manager = object.__new__(ReqManagerForMamba) + manager.recurrent_state = object.__new__(RecurrentState) config = SimpleNamespace( linear_layer_num=2, conv_state_dtype=torch.bfloat16, @@ -509,23 +511,24 @@ def test_checkpoint_restores_pending_history_and_accepted_conv(save_big, monkeyp conv_kernel_size=4, ) manager.linear_config = config + manager.recurrent_state.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, dtype, (2, 32, 64), 2, "cuda") - manager.req_to_ssm_state.buffer.zero_() - manager.ssm_update_cache = ReplaySSMCache(manager.req_to_ssm_state.buffer, 16, 3) - cache = manager.ssm_update_cache + manager.recurrent_state.slots_per_req = 1 + manager.recurrent_state.accepted_index = torch.zeros(4, device="cuda", dtype=torch.int32) + manager.recurrent_state.accepted_index[1] = 1 + manager.recurrent_state.conv_state = LayerCache(4, torch.bfloat16, (192, 5), 2, "cuda") + manager.recurrent_state.conv_state.buffer.copy_(torch.randn_like(manager.recurrent_state.conv_state.buffer)) + manager.recurrent_state.ssm_state = LayerCache(4, dtype, (2, 32, 64), 2, "cuda") + manager.recurrent_state.ssm_state.buffer.zero_() + manager.recurrent_state.updates = ReplaySSMCache(manager.recurrent_state.ssm_state.buffer, 16, 3) + cache = manager.recurrent_state.updates 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 - active_before = manager.req_to_ssm_state.buffer[:, 1].clone() - expected = manager.req_to_ssm_state.buffer[:, 1].float().clone() + active_before = manager.recurrent_state.ssm_state.buffer[:, 1].clone() + expected = manager.recurrent_state.ssm_state.buffer[:, 1].float().clone() for i in range(2): key = cache.raw_keys[:, 1, :, i].float() key /= (key.square().sum(-1, keepdim=True) + 1e-6).sqrt() @@ -535,7 +538,7 @@ def test_checkpoint_restores_pending_history_and_accepted_conv(save_big, monkeyp ) expected += key[..., None] * delta[..., None, :] expected = expected.to(dtype) - conv_expected = manager.req_to_conv_state.buffer[:, 1, :, 1:4].clone() + conv_expected = manager.recurrent_state.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) @@ -551,25 +554,21 @@ def test_checkpoint_restores_pending_history_and_accepted_conv(save_big, monkeyp # a tolerance for FP32 reduction differences crossing BF16 rounding points. torch.testing.assert_close(state_saved, materialized.cpu(), rtol=0, atol=0) torch.testing.assert_close(conv_saved, conv_expected.cpu(), rtol=0, atol=0) - torch.testing.assert_close(manager.req_to_ssm_state.buffer[:, 1], active_before, rtol=0, atol=0) + torch.testing.assert_close(manager.recurrent_state.ssm_state.buffer[:, 1], active_before, rtol=0, atol=0) assert cache.cursors[1].item() == 2 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], materialized, rtol=0, atol=0) - torch.testing.assert_close(manager.req_to_conv_state.buffer[:, 2, :, :3], conv_expected, rtol=0, atol=0) + assert manager.recurrent_state.accepted_index[2].item() == 0 + torch.testing.assert_close(manager.recurrent_state.ssm_state.buffer[:, 2], materialized, rtol=0, atol=0) + torch.testing.assert_close(manager.recurrent_state.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, - ssm_update_cache=cache, - ssm_slots_per_req=1, - req_to_mtp_state_index=manager.req_to_mtp_state_index, + recurrent_state=manager.recurrent_state, ) helper = Qwen3NextLinearAttPageHelper(mem) # A second pending suffix exercises PD export, independently of CPU save. @@ -591,17 +590,17 @@ def test_checkpoint_restores_pending_history_and_accepted_conv(save_big, monkeyp ssm_page = torch.empty(helper.ssm_shape, device="cuda", dtype=dtype) helper._write_one_rank(mem, 0, 1, conv_page, ssm_page) torch.testing.assert_close(ssm_page, expected, rtol=state_rtol, atol=2e-5) - torch.testing.assert_close(manager.req_to_ssm_state.buffer[:, 1], active_before, rtol=0, atol=0) + torch.testing.assert_close(manager.recurrent_state.ssm_state.buffer[:, 1], active_before, rtol=0, atol=0) assert cache.cursors[1].item() == 1 cache.cursors[2] = 7 - manager.req_to_mtp_state_index[2] = 2 + manager.recurrent_state.accepted_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], ssm_page, rtol=0, atol=0) - 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 + torch.testing.assert_close(manager.recurrent_state.ssm_state.buffer[:, 2], ssm_page, rtol=0, atol=0) + torch.testing.assert_close(manager.recurrent_state.conv_state.buffer[:, 2, :, :3], conv_expected, rtol=0, atol=0) + assert cache.cursors[2].item() == manager.recurrent_state.accepted_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 + assert torch.count_nonzero(manager.recurrent_state.ssm_state.buffer[:, 2]).item() == 0 + assert torch.count_nonzero(manager.recurrent_state.conv_state.buffer[:, 2]).item() == 0 @pytest.mark.parametrize("dtype", [torch.float32, torch.bfloat16]) diff --git a/unit_tests/common/basemodel/triton_kernel/linear_att/test_replayssm_execution.py b/unit_tests/common/basemodel/triton_kernel/linear_att/test_replayssm_execution.py index 1b6c479978..e2ba53add1 100644 --- a/unit_tests/common/basemodel/triton_kernel/linear_att/test_replayssm_execution.py +++ b/unit_tests/common/basemodel/triton_kernel/linear_att/test_replayssm_execution.py @@ -5,6 +5,7 @@ from lightllm.common.basemodel.triton_kernel.linear_att.mtp_state_params import ( build_dynamic_mtp_linear_att_state_params, ) +from lightllm.common.recurrent_state import RecurrentState from lightllm.common.req_manager.linear_att import ReqManagerForMamba pytestmark = pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") @@ -32,8 +33,9 @@ def test_request_manager_accept_with_replay_graph(dtype, capture_metadata, varle for cls in [ReplaySSMCache, ReplaySSMCache] ] manager = object.__new__(ReqManagerForMamba) - manager.ssm_update_cache = caches[1] - manager.req_to_mtp_state_index = torch.zeros(5, device="cuda", dtype=torch.int32) + manager.recurrent_state = object.__new__(RecurrentState) + manager.recurrent_state.updates = caches[1] + manager.recurrent_state.accepted_index = torch.zeros(5, device="cuda", dtype=torch.int32) rows = torch.full((12,), 4, device="cuda", dtype=torch.int32) mtp = torch.zeros_like(rows) flags = torch.zeros_like(rows) @@ -71,7 +73,7 @@ def inputs(iteration): return ids, lengths, counts, logical def metadata(cache): - cu, reqs, _ = build_dynamic_mtp_linear_att_state_params(rows, mtp, manager.req_to_mtp_state_index, 4) + cu, reqs, _ = build_dynamic_mtp_linear_att_state_params(rows, mtp, manager.recurrent_state.accepted_index, 4) return cu, reqs, cache.prepare_decode(reqs, cu) def forward_accept(index, cu, reqs): @@ -90,7 +92,7 @@ def forward_accept(index, cu, reqs): forward_accept(index, cu, reqs) cache.state.copy_(initial) cache.cursors.zero_() - manager.req_to_mtp_state_index.zero_() + manager.recurrent_state.accepted_index.zero_() g = torch.cuda.CUDAGraph() with torch.cuda.graph(g): if capture_metadata: @@ -102,7 +104,7 @@ def forward_accept(index, cu, reqs): for cache in caches: cache.state.copy_(initial) cache.cursors.zero_() - manager.req_to_mtp_state_index.zero_() + manager.recurrent_state.accepted_index.zero_() folds = 0 for iteration in range(32): ids, lengths, counts, logical = inputs(iteration) @@ -126,7 +128,7 @@ def forward_accept(index, cu, reqs): torch.testing.assert_close( caches[1].state.float(), caches[0].state.float(), atol=0.002, rtol=torch.finfo(dtype).eps ) - table = manager.req_to_mtp_state_index.cpu().tolist() + table = manager.recurrent_state.accepted_index.cpu().tolist() for req, count in zip(ids, counts): assert table[req] == count - 1 assert table[2] == table[4] == 0 diff --git a/unit_tests/common/basemodel/triton_kernel/linear_att/test_ssm_autotune.py b/unit_tests/common/basemodel/triton_kernel/linear_att/test_ssm_autotune.py index 36f550a04c..b8e6329d68 100644 --- a/unit_tests/common/basemodel/triton_kernel/linear_att/test_ssm_autotune.py +++ b/unit_tests/common/basemodel/triton_kernel/linear_att/test_ssm_autotune.py @@ -7,6 +7,7 @@ from lightllm.common.basemodel.triton_kernel.linear_att import ssm_autotune as tuning from lightllm.common.basemodel.triton_kernel.linear_att.replayssm import ReplaySSMCache from lightllm.common.basemodel.triton_kernel.linear_att.replayssm_compact import CompactSSMCache +from lightllm.common.recurrent_state import RecurrentState, RecurrentStateContext from lightllm.common.kernel_config import KernelConfigs from lightllm.common import kernel_config as kernel_config_module from lightllm.common.triton_utils import autotuner as autotuner_module @@ -318,8 +319,9 @@ def test_two_verify_graphs_accept_with_their_own_configs(monkeypatch, mode, proj configs = [{"BV": 8, "num_warps": 1, "num_stages": 1}, {"BV": 32, "num_warps": 2, "num_stages": 1}] monkeypatch.setattr(tuning, "select_config", lambda *args: configs[args[2].shape[1] > 8]) manager = object.__new__(ReqManagerForMamba) - manager.ssm_update_cache = cache - manager.req_to_mtp_state_index = torch.zeros(9, dtype=torch.int32, device="cuda") + manager.recurrent_state = object.__new__(RecurrentState) + manager.recurrent_state.updates = cache + manager.recurrent_state.accepted_index = torch.zeros(9, dtype=torch.int32, device="cuda") buckets = [] for index, ids in enumerate(([0], [1, 3])): # The physical verify buckets include HOLD, but acceptance is unpadded. @@ -349,14 +351,14 @@ def forward(target): state = SimpleNamespace( hidden_collector=SimpleNamespace(finish_output=lambda **kw: ModelMtpOutputCollector()), prompt_logics=None, - ssm_run_config=config, + recurrent_state_context=RecurrentStateContext(config), logits_row_stride=1, ) output = TpPartBaseModel._create_model_output( None, PostLayerOutput(logits=outputs[0].view(count * 4, -1)), state ) output = TpPartBaseModel._create_unpad_decode_model_output(None, output, len(ids) * 4) - assert output.ssm_run_config is config + assert output.recurrent_state_context._run_config is config buckets.append((graph, outputs, output, ref, args, reqs, cu, ids)) cache.state.copy_(initial) if mode == "replay": @@ -382,10 +384,10 @@ def forward(target): mtp, flags, 4, - ssm_accept_batches=tuple((len(b[-1]), b[2].ssm_run_config) for b in buckets), + recurrent_state_batches=tuple((len(b[-1]), b[2].recurrent_state_context) for b in buckets), ) for _, _, _, ref, _, reqs, _, ids in buckets: - ref.accept_updates(reqs[:-1], manager.req_to_mtp_state_index) + ref.accept_updates(reqs[:-1], manager.recurrent_state.accepted_index) torch.testing.assert_close(cache.state[:, ids], ref.state[:, ids], atol=0, rtol=0) if mode == "replay": assert torch.equal(cache.cursors[ids], ref.cursors[ids]) diff --git a/unit_tests/common/test_recurrent_state.py b/unit_tests/common/test_recurrent_state.py new file mode 100644 index 0000000000..227548e9b6 --- /dev/null +++ b/unit_tests/common/test_recurrent_state.py @@ -0,0 +1,112 @@ +from types import SimpleNamespace + +import pytest +import torch + +from lightllm.common.kv_cache_mem_manager.qwen3next_mem_manager import Qwen3NextLinearAttPageHelper +from lightllm.common.recurrent_state import RecurrentState +from lightllm.common.req_manager.linear_att import ReqManagerForMamba +from lightllm.common.state_cache_manager import LinearAttCacheConfig, LinearAttCacheManager + + +pytestmark = pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") + + +@pytest.mark.parametrize("mode,step", [("native", 0), ("native", 2), ("compact", 2), ("replay", 0), ("replay", 2)]) +@pytest.mark.parametrize("dtype", [torch.float32, torch.bfloat16]) +def test_checkpoint_and_pd_export_same_accepted_prefix_and_restore_reused_slot(mode, step, dtype): + config = LinearAttCacheConfig( + tp_world_size=1, + full_att_all_num_kv_heads=1, + full_att_dtype=torch.bfloat16, + full_att_num_kv_heads=1, + full_att_head_dim=32, + global_linear_k_heads=1, + global_linear_v_heads=2, + num_linear_k_heads=1, + num_linear_v_heads=2, + head_linear_k_dim=32, + head_linear_v_dim=32, + conv_kernel_size=4, + linear_layer_num=2, + conv_state_dtype=torch.bfloat16, + ssm_state_dtype=dtype, + full_attention_interval=3, + all_layer_num=3, + ) + args = SimpleNamespace(mtp_step=step, ssm_state_mode=mode, ssm_replay_config={"cache_len": 4}) + state = RecurrentState(3, config, args) + manager = object.__new__(ReqManagerForMamba) + manager.recurrent_state = state + cache = LinearAttCacheManager(3, config) + manager.mem_manager = SimpleNamespace(big_page_buffers=cache) + mem = SimpleNamespace(linear_config=config, recurrent_state=state) + helper = Qwen3NextLinearAttPageHelper(mem) + src, dst, offset = 1, 2, min(1, step) + width = config.get_mtp_conv_state_shape(step)[-1] + state.conv_state.buffer.copy_(torch.arange(width, device="cuda").to(torch.bfloat16)) + for slot in range(state.ssm_state.buffer.shape[1]): + state.ssm_state.buffer[:, slot].fill_(slot + 1) + if state.accepted_index is not None: + state.accepted_index[src] = offset + expected_conv = state.conv_state.buffer[:, src, ..., offset : offset + 3].clone() + expected_ssm = state.ssm_state.buffer[:, src * state.slots_per_req + (offset if mode == "native" else 0)].clone() + if mode == "replay": + # One pending accepted update: beta=0 leaves only decay of the base. + state.updates.raw_keys.zero_() + state.updates.raw_values.zero_() + state.updates.gates.fill_(-0.5) + state.updates.betas.zero_() + state.updates.cursors[src] = 1 + expected_ssm = (expected_ssm.float() * torch.exp(torch.tensor(-0.5))).to(dtype) + before = state.ssm_state.buffer.clone() + snapshot = state.snapshot(src) + torch.testing.assert_close(snapshot.ssm, expected_ssm, rtol=2e-5, atol=2e-5) + expected_ssm = snapshot.ssm.clone() + + manager.save_state(src, 0, cache) + manager.save_big_page_states(torch.tensor([src], dtype=torch.int32, device="cuda"), [src], [1]) + conv_page = torch.empty(helper.conv_shape, device="cuda", dtype=torch.bfloat16) + ssm_page = torch.empty(helper.ssm_shape, device="cuda", dtype=dtype) + helper._write_one_rank(mem, 0, src, conv_page, ssm_page) + torch.cuda.synchronize() + for index in (0, 1): + conv, ssm = cache.get_state_cache(index) + torch.testing.assert_close(conv, expected_conv.cpu(), rtol=0, atol=0) + torch.testing.assert_close(ssm, expected_ssm.cpu(), rtol=0, atol=0) + torch.testing.assert_close(conv_page, expected_conv, rtol=0, atol=0) + torch.testing.assert_close(ssm_page, expected_ssm, rtol=0, atol=0) + assert torch.equal(state.ssm_state.buffer, before) + if mode == "replay": + assert state.updates.cursors[src].item() == 1 + + def dirty_destination(): + if state.accepted_index is not None: + state.accepted_index[dst] = step + if mode == "replay": + state.updates.cursors[dst] = 3 + + for restore in ( + lambda: manager.restore_state(SimpleNamespace(req_idx=dst), cache, 0), + lambda: helper._read_one_rank(mem, 0, dst, conv_page, ssm_page), + ): + dirty_destination() + restore() + actual = state.snapshot(dst) + torch.testing.assert_close(actual.conv, expected_conv, rtol=0, atol=0) + torch.testing.assert_close(actual.ssm, expected_ssm, rtol=0, atol=0) + if state.accepted_index is not None: + assert state.accepted_index[dst].item() == 0 + if mode == "replay": + assert state.updates.cursors[dst].item() == 0 + + # A skipped checkpoint must not overwrite another cache owner. + untouched = tuple(x.clone() for x in cache.get_state_cache(2)) + manager.save_big_page_states(torch.tensor([src], dtype=torch.int32, device="cuda"), [src], [-1]) + torch.cuda.synchronize() + for actual, expected in zip(cache.get_state_cache(2), untouched): + assert torch.equal(actual, expected) + state.reset(dst) + assert not state.conv_state.buffer[:, dst].any() + start = dst * state.slots_per_req + assert not state.ssm_state.buffer[:, start : start + state.slots_per_req].any() diff --git a/unit_tests/models/glm5_next/test_cache.py b/unit_tests/models/glm5_next/test_cache.py index ce461914f6..482b564001 100644 --- a/unit_tests/models/glm5_next/test_cache.py +++ b/unit_tests/models/glm5_next/test_cache.py @@ -98,10 +98,10 @@ def test_hybrid_checkpoint_restore_and_packed_kv_copy(monkeypatch, small_page, t slot = cache.alloc_one_state_cache() source_req = SimpleNamespace(req_idx=0) req.init_hybrid_attention_state(source_req) - req.req_to_conv_state.buffer[:, 0].normal_() - req.req_to_ssm_state.buffer[:, 0].normal_() - conv = req.req_to_conv_state.buffer[:, 0, ..., :3].clone() - updates = req.ssm_update_cache + req.recurrent_state.conv_state.buffer[:, 0].normal_() + req.recurrent_state.ssm_state.buffer[:, 0].normal_() + conv = req.recurrent_state.conv_state.buffer[:, 0, ..., :3].clone() + updates = req.recurrent_state.updates if updates is not None: # Save an accepted suffix, not merely an already folded checkpoint. heads, dim = config.num_linear_v_heads, config.head_linear_k_dim @@ -120,16 +120,16 @@ def test_hybrid_checkpoint_restore_and_packed_kv_copy(monkeypatch, small_page, t updates.accept_updates(ids, accepted) ssm = updates.snapshot_accepted_state(0).clone() if mtp_step: - req.req_to_mtp_state_index[2] = mtp_step + req.recurrent_state.accepted_index[2] = mtp_step else: - ssm = req.req_to_ssm_state.buffer[:, 0].clone() + ssm = req.recurrent_state.ssm_state.buffer[:, 0].clone() if small_page: req.save_state(0, slot, cache) else: req.save_big_page_states(torch.tensor([0], dtype=torch.int32, device="cuda"), [0], [slot]) torch.cuda.synchronize() - req.req_to_conv_state.buffer[:, 0].zero_() - req.req_to_ssm_state.buffer[:, 0].zero_() + req.recurrent_state.conv_state.buffer[:, 0].zero_() + req.recurrent_state.ssm_state.buffer[:, 0].zero_() req.req_to_indexer_tail.buffer.normal_() unchanged = req.req_to_indexer_tail.buffer[:, [0, 1, 3]].clone() dest_req = SimpleNamespace(req_idx=2, shared_kv_node=SimpleNamespace(small_page_buffer_idx=slot)) @@ -143,17 +143,19 @@ def test_hybrid_checkpoint_restore_and_packed_kv_copy(monkeypatch, small_page, t if mode == "replay": assert updates.cursors[2] == 0 if mtp_step: - assert req.req_to_mtp_state_index[2] == 0 - assert torch.equal(req.req_to_conv_state.buffer[:, 2, ..., :3], conv) - assert torch.equal(req.req_to_ssm_state.buffer[:, 2 * req.ssm_slots_per_req], ssm) + assert req.recurrent_state.accepted_index[2] == 0 + assert torch.equal(req.recurrent_state.conv_state.buffer[:, 2, ..., :3], conv) + assert torch.equal(req.recurrent_state.ssm_state.buffer[:, 2 * req.recurrent_state.slots_per_req], ssm) # Both page sizes are aligned to complete pools. Restoring a prefix must # discard stale tail values from a previously allocated request slot. assert not req.req_to_indexer_tail.buffer[:, 2].any() assert torch.equal(req.req_to_indexer_tail.buffer[:, [0, 1, 3]], unchanged) req.req_to_indexer_tail.buffer[:, 2].normal_() req.init_hybrid_attention_state(dest_req) - assert not req.req_to_conv_state.buffer[:, 2].any() - assert not req.req_to_ssm_state.buffer[:, 2 * req.ssm_slots_per_req : 3 * req.ssm_slots_per_req].any() + assert not req.recurrent_state.conv_state.buffer[:, 2].any() + assert not req.recurrent_state.ssm_state.buffer[ + :, 2 * req.recurrent_state.slots_per_req : 3 * req.recurrent_state.slots_per_req + ].any() assert not req.req_to_indexer_tail.buffer[:, 2].any() assert torch.equal(req.req_to_indexer_tail.buffer[:, [0, 1, 3]], unchanged) # KV moves carry MLA latents and pooled FP8 bytes; raw keys/gates only diff --git a/unit_tests/models/glm5_next/test_mtp.py b/unit_tests/models/glm5_next/test_mtp.py index 0fe7514fc7..1bd3145077 100644 --- a/unit_tests/models/glm5_next/test_mtp.py +++ b/unit_tests/models/glm5_next/test_mtp.py @@ -1,6 +1,7 @@ import pytest import torch +from lightllm.common.recurrent_state import RecurrentState from lightllm.common.basemodel.triton_kernel.linear_att.fla.ops.kda_decode import fused_recurrent_kda from lightllm.models.glm5_next.triton_kernel.kpool import compress_pools @@ -190,12 +191,14 @@ def test_kda_backend_routes_accepted_conv_and_ssm_states(dynamic, mode, mtp_step uses_dynamic_spec_verify_layout=lambda: dynamic, split_qkv=lambda x: x.split(hidden, -1), ) + recurrent_state = object.__new__(RecurrentState) + recurrent_state.accepted_index = accepted_offsets + recurrent_state.updates = None + recurrent_state.slots_per_req = width manager = SimpleNamespace( - req_to_mtp_state_index=accepted_offsets, + recurrent_state=recurrent_state, HOLD_REQUEST_ID=3, get_mamba_cache=lambda layer: (conv, ssm), - ssm_update_cache=None, - ssm_slots_per_req=width, linear_config=SimpleNamespace(full_attention_interval=4), ) if mode != "native": @@ -205,11 +208,11 @@ def test_kda_backend_routes_accepted_conv_and_ssm_states(dynamic, mode, mtp_step indexes = torch.arange(4, device="cuda") * width + accepted_offsets ssm = ssm[indexes].clone() before_verify = ssm.clone() - manager.ssm_slots_per_req = 1 + manager.recurrent_state.slots_per_req = 1 if mode == "replay": - manager.ssm_update_cache = ReplaySSMCache(ssm.unsqueeze(0), 8, width, kda=True) + manager.recurrent_state.updates = ReplaySSMCache(ssm.unsqueeze(0), 8, width, kda=True) else: - manager.ssm_update_cache = CompactSSMCache(ssm.unsqueeze(0), width, torch.bfloat16, kda=True) + manager.recurrent_state.updates = CompactSSMCache(ssm.unsqueeze(0), width, torch.bfloat16, kda=True) infer = SimpleNamespace( batch_size=tokens, b_req_idx=torch.tensor(request_rows, device="cuda", dtype=torch.int32), @@ -243,8 +246,8 @@ def test_kda_backend_routes_accepted_conv_and_ssm_states(dynamic, mode, mtp_step for req, length in zip(requests, lengths): accepted[req] = min(length, 2) - 1 if mtp_step: - manager.ssm_update_cache.accept_updates(state.b_conv_buffer_idx, accepted) - manager.ssm_update_cache.merge_accepted_updates(state.b_conv_buffer_idx) + manager.recurrent_state.updates.accept_updates(state.b_conv_buffer_idx, accepted) + manager.recurrent_state.updates.merge_accepted_updates(state.b_conv_buffer_idx) for req in requests: torch.testing.assert_close(ssm[req], expected_ssm[req * width + accepted[req]], atol=2e-3, rtol=1e-2) if dynamic and mtp_step: @@ -309,9 +312,10 @@ def test_kda_prefill_materializes_accepted_history_and_uses_single_state_slot(mo (None, native_conv, native_ssm, width), (updates, conv, checkpoint[0], 1), ]: - manager = SimpleNamespace( - ssm_update_cache=cache, ssm_slots_per_req=slots, get_mamba_cache=lambda layer: (conv_pool, ssm_pool) - ) + recurrent_state = object.__new__(RecurrentState) + recurrent_state.updates = cache + recurrent_state.slots_per_req = slots + manager = SimpleNamespace(recurrent_state=recurrent_state, get_mamba_cache=lambda layer: (conv_pool, ssm_pool)) infer = SimpleNamespace( b_req_idx=ids, b_ready_cache_len=ids.new_tensor([2]), diff --git a/unit_tests/models/glm5_next/test_pd_cache.py b/unit_tests/models/glm5_next/test_pd_cache.py index 4983679b69..8d7300a84b 100644 --- a/unit_tests/models/glm5_next/test_pd_cache.py +++ b/unit_tests/models/glm5_next/test_pd_cache.py @@ -33,8 +33,8 @@ def make_mems(monkeypatch, request): def check_shm_refs(mem, req): assert mem.req_to_indexer_tail is req.req_to_indexer_tail - assert mem.req_to_conv_state is req.req_to_conv_state - assert mem.req_to_ssm_state is req.req_to_ssm_state + assert mem.recurrent_state.conv_state is req.recurrent_state.conv_state + assert mem.recurrent_state.ssm_state is req.recurrent_state.ssm_state assert mem.big_page_buffers is None monkeypatch.setattr(MemoryManager, "write_to_shm", check_shm_refs) @@ -67,8 +67,8 @@ def create(tp, mode=None): for _ in range(tp): mem = Glm5NextMemManager(32, torch.bfloat16, 1, 584, 2 + int(mtp_step > 0), config) req = Glm5NextReqManager(3, 32, mem, config) - assert req.ssm_slots_per_req == (mtp_step + 1 if args.ssm_state_mode == "native" else 1) - assert (req.ssm_update_cache is None) == (args.ssm_state_mode == "native") + assert req.recurrent_state.slots_per_req == (mtp_step + 1 if args.ssm_state_mode == "native" else 1) + assert (req.recurrent_state.updates is None) == (args.ssm_state_mode == "native") mem.write_to_shm(req) assert mem.big_page_buffers is not None mems.append(mem) @@ -98,10 +98,10 @@ def test_pd_kv_and_runtime_state_roundtrip(make_mems, prefill_tp, decode_tp, rem for rank, mem in enumerate(source): heads = slice(rank * 8 // prefill_tp, (rank + 1) * 8 // prefill_tp) mem.kv_buffer.view(torch.uint8)[:, src_indexes] = packed - mem.req_to_conv_state.buffer[:, src_req, ..., :3].copy_(conv[:, :, heads].reshape(6, -1, 3)) - mem.req_to_ssm_state.buffer[:, src_req * mem.ssm_slots_per_req].copy_(ssm[:, heads]) + mem.recurrent_state.conv_state.buffer[:, src_req, ..., :3].copy_(conv[:, :, heads].reshape(6, -1, 3)) + mem.recurrent_state.ssm_state.buffer[:, src_req * mem.recurrent_state.slots_per_req].copy_(ssm[:, heads]) mem.req_to_indexer_tail.buffer[:, src_req].copy_(tail) - updates = mem.ssm_update_cache + updates = mem.recurrent_state.updates if updates is not None: # The exported checkpoint must contain accepted, unfurled KDA history. h = mem.linear_config.num_linear_v_heads @@ -130,19 +130,19 @@ def test_pd_kv_and_runtime_state_roundtrip(make_mems, prefill_tp, decode_tp, rem untouched = [] for mem in dest: mem.kv_buffer.zero_() - mem.req_to_conv_state.buffer.normal_() - mem.req_to_ssm_state.buffer.normal_() + mem.recurrent_state.conv_state.buffer.normal_() + mem.recurrent_state.ssm_state.buffer.normal_() mem.req_to_indexer_tail.buffer.normal_() - if mem.ssm_update_cache is not None and hasattr(mem.ssm_update_cache, "cursors"): - mem.ssm_update_cache.cursors[dst_req] = 5 - if mem.req_to_mtp_state_index is not None: - mem.req_to_mtp_state_index[dst_req] = mtp_size - 1 + if mem.recurrent_state.updates is not None and hasattr(mem.recurrent_state.updates, "cursors"): + mem.recurrent_state.updates.cursors[dst_req] = 5 + if mem.recurrent_state.accepted_index is not None: + mem.recurrent_state.accepted_index[dst_req] = mtp_size - 1 untouched.append( [ x.buffer[:, torch.tensor([0, 1, 3], device="cuda") * stride].clone() for x, stride in ( - (mem.req_to_conv_state, 1), - (mem.req_to_ssm_state, mem.ssm_slots_per_req), + (mem.recurrent_state.conv_state, 1), + (mem.recurrent_state.ssm_state, mem.recurrent_state.slots_per_req), (mem.req_to_indexer_tail, 1), ) ] @@ -160,15 +160,23 @@ def test_pd_kv_and_runtime_state_roundtrip(make_mems, prefill_tp, decode_tp, rem for rank, mem in enumerate(dest): heads = slice(rank * 8 // decode_tp, (rank + 1) * 8 // decode_tp) assert torch.equal(mem.kv_buffer.view(torch.uint8)[:, dst_indexes], packed) - assert torch.equal(mem.req_to_conv_state.buffer[:, dst_req, ..., :3], conv[:, :, heads].reshape(6, -1, 3)) - assert torch.equal(mem.req_to_ssm_state.buffer[:, dst_req * mem.ssm_slots_per_req], ssm[:, heads]) + assert torch.equal( + mem.recurrent_state.conv_state.buffer[:, dst_req, ..., :3], conv[:, :, heads].reshape(6, -1, 3) + ) + assert torch.equal( + mem.recurrent_state.ssm_state.buffer[:, dst_req * mem.recurrent_state.slots_per_req], ssm[:, heads] + ) assert torch.equal(mem.req_to_indexer_tail.buffer[:, dst_req], tail) - if mem.ssm_update_cache is not None and hasattr(mem.ssm_update_cache, "cursors"): - assert mem.ssm_update_cache.cursors[dst_req] == 0 - if mem.req_to_mtp_state_index is not None and mem.ssm_update_cache is not None: - assert mem.req_to_mtp_state_index[dst_req] == 0 + if mem.recurrent_state.updates is not None and hasattr(mem.recurrent_state.updates, "cursors"): + assert mem.recurrent_state.updates.cursors[dst_req] == 0 + if mem.recurrent_state.accepted_index is not None and mem.recurrent_state.updates is not None: + assert mem.recurrent_state.accepted_index[dst_req] == 0 for (state, stride), expected in zip( - ((mem.req_to_conv_state, 1), (mem.req_to_ssm_state, mem.ssm_slots_per_req), (mem.req_to_indexer_tail, 1)), + ( + (mem.recurrent_state.conv_state, 1), + (mem.recurrent_state.ssm_state, mem.recurrent_state.slots_per_req), + (mem.req_to_indexer_tail, 1), + ), untouched[rank], ): indexes = torch.tensor([0, 1, 3], device="cuda") * stride diff --git a/unit_tests/models/test_post_layer_output.py b/unit_tests/models/test_post_layer_output.py index c8e081862d..43ff69e8c3 100644 --- a/unit_tests/models/test_post_layer_output.py +++ b/unit_tests/models/test_post_layer_output.py @@ -291,7 +291,7 @@ def test_model_output_and_unpadding_preserve_candidate_mapping() -> None: logits = torch.randn(3, 2) ids = torch.tensor([[10, 20], [30, 40], [50, 60]]) state = SimpleNamespace( - hidden_collector=NoopHiddenCollector(), prompt_logics=None, ssm_run_config=None, logits_row_stride=1 + hidden_collector=NoopHiddenCollector(), prompt_logics=None, recurrent_state_context=None, logits_row_stride=1 ) output = model._create_model_output(PostLayerOutput(logits, ids), state) unpadded = model._create_unpad_decode_model_output(output, 2) From 92127025c5417dcfc5c1ba312a3551538038540a Mon Sep 17 00:00:00 2001 From: sufubao Date: Sat, 17 Oct 2026 17:16:39 +0800 Subject: [PATCH 20/40] fix: retain FP32 replay history for FP16 activations --- .../triton_kernel/linear_att/replayssm.py | 8 ++- .../linear_att/test_replayssm.py | 57 +++++++++++++++++++ .../triton_kernel/test_greedy_sampling.py | 1 + 3 files changed, 65 insertions(+), 1 deletion(-) diff --git a/lightllm/common/basemodel/triton_kernel/linear_att/replayssm.py b/lightllm/common/basemodel/triton_kernel/linear_att/replayssm.py index 3db7cf2f98..c1f4ca271c 100644 --- a/lightllm/common/basemodel/triton_kernel/linear_att/replayssm.py +++ b/lightllm/common/basemodel/triton_kernel/linear_att/replayssm.py @@ -627,9 +627,15 @@ def __init__( self.forward_cursors = torch.zeros_like(self.cursors) # Alternate halves on fold so CTAs cannot overwrite history another # V tile is still reading. Cursor packs phase (2*L) and count (0..L). + # Derived corrections can exceed the activation/state range even when + # the recurrence stays finite. Only BF16 has enough exponent range for + # the short-history narrowing optimization; FP16 must retain FP32. history_dtype = ( activation_dtype - if verify_width > 1 and projection_mode == "precompute" and capacity <= 8 + if activation_dtype == torch.bfloat16 + and verify_width > 1 + and projection_mode == "precompute" + and capacity <= 8 else torch.float32 ) key_shape = (layers, slots, self.num_key_heads, 2 * capacity, k) 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 index bae5eba2cd..5bcacc8ecf 100644 --- a/unit_tests/common/basemodel/triton_kernel/linear_att/test_replayssm.py +++ b/unit_tests/common/basemodel/triton_kernel/linear_att/test_replayssm.py @@ -24,6 +24,63 @@ def reference(q, k, v, a, b, a_log, bias, state, kda=False, lower_bound=-5.0): return torch.einsum("...kv,...k->...v", state, q), state +@pytest.mark.parametrize("capacity", [4, 8]) +@pytest.mark.parametrize("bv", [8, 32], ids=["small_bv_fallback", "precomputed"]) +@pytest.mark.parametrize("state_dtype", [torch.float32, torch.bfloat16]) +@pytest.mark.parametrize("kda", [False, True], ids=["gdn", "kda"]) +def test_fp16_correction_history_preserves_finite_recurrence(capacity, bv, state_dtype, kda): + state = torch.zeros(1, 2, 1, 128, 128, dtype=state_dtype, device="cuda") + cache = ReplaySSMCache( + state, + capacity, + 2, + torch.float16, + num_key_heads=1, + kda=kda, + projection_mode="precompute", + run_config={"BV": bv, "num_warps": 1}, + ) + reqs = torch.tensor([0, 1], dtype=torch.int32, device="cuda") + cu = torch.tensor([0, 2, 4], dtype=torch.int32, device="cuda") + accepted = torch.tensor([1, 0], dtype=torch.int32, device="cuda") + q = torch.zeros(1, 4, 1, 128, dtype=torch.float16, device="cuda") + q[0, :2, 0, 0] = 1 + k = torch.zeros_like(q) + k[0, 0, 0, 0] = -0.25 + k[0, 1, 0, 0] = 1 + v = torch.zeros_like(q) + v[0, :2, 0, 0] = 60000 + a = torch.full((4, 128 if kda else 1), -20, dtype=torch.float16, device="cuda") + b = torch.full((4, 1), 20, dtype=torch.float16, device="cuda") + alog = torch.zeros(1, device="cuda") + bias = torch.zeros(128 if kda else 1, device="cuda") + expected_state = state[0, 0].float().clone() + for round_idx in range(2): + cache.prepare_decode(reqs, cu) + # The second round must consume the first round's accepted history. + # For capacity 4 the next span exactly fills it, so no fold should hide + # an overflowed derived correction by reconstructing from raw inputs. + assert cache.cursors[0].item() == 2 * round_idx + assert not state.any() + actual = cache.forward(0, q, k, v, a, b, alog, bias, reqs, cu) + expected = [] + for token in range(2): + output, expected_state = reference( + q[0, token], k[0, token], v[0, token], a[token], b[token], alog, bias, expected_state, kda=kda + ) + expected.append(output) + assert torch.isfinite(expected_state).all() + assert torch.isfinite(actual).all() + torch.testing.assert_close(actual[0, :2].float(), torch.stack(expected), atol=4, rtol=1e-3) + cache.accept_updates(reqs, accepted) + # The correction is larger than FP16 can represent even though the + # recurrent state, inputs, and attention outputs are all finite. + assert cache.deltas[0, 0, 0, 1, 0].item() > torch.finfo(torch.float16).max + assert torch.isfinite(cache.deltas[0, 0, :, :4]).all() + assert cache.cursors[1].item() == 0 + torch.testing.assert_close(cache.snapshot_accepted_state(0)[0].float(), expected_state, atol=128, rtol=0.01) + + @pytest.mark.parametrize("capacity", [16, 32, 64]) @pytest.mark.parametrize("width", [1, 3, 5, 16]) @pytest.mark.parametrize("dims", [(32, 64), (128, 128)]) diff --git a/unit_tests/common/basemodel/triton_kernel/test_greedy_sampling.py b/unit_tests/common/basemodel/triton_kernel/test_greedy_sampling.py index ee7980b3a2..7ee34afe72 100644 --- a/unit_tests/common/basemodel/triton_kernel/test_greedy_sampling.py +++ b/unit_tests/common/basemodel/triton_kernel/test_greedy_sampling.py @@ -65,6 +65,7 @@ def test_full_sampler_with_penalties_invalid_ids_and_random_sampling(monkeypatch vocab, batch = 4101, 4 manager = SimpleNamespace( penalty_counter_mode="gpu_counter", + req_to_bitmask_enabled=None, vocab_size=vocab, req_to_presence_penalty=torch.full((batch,), 0.0 if identity else 0.5, device="cuda"), req_to_frequency_penalty=torch.full((batch,), 0.0 if identity else 0.1, device="cuda"), From cc9d9a5f48b923a511a2791d524ef8e81ddd91e1 Mon Sep 17 00:00:00 2001 From: sufubao Date: Sat, 17 Oct 2026 17:17:14 +0800 Subject: [PATCH 21/40] fix: run Triton fallback on the failing FA3 invocation --- .../vit/triton_kernel/flashattention_nopad.py | 4 +- .../vit/test_flash_attention_forward.py | 52 +++++++++++++++++++ 2 files changed, 54 insertions(+), 2 deletions(-) diff --git a/lightllm/models/vit/triton_kernel/flashattention_nopad.py b/lightllm/models/vit/triton_kernel/flashattention_nopad.py index a9cbcda71d..af062c2a5b 100644 --- a/lightllm/models/vit/triton_kernel/flashattention_nopad.py +++ b/lightllm/models/vit/triton_kernel/flashattention_nopad.py @@ -222,8 +222,8 @@ def flash_attention_fwd(q, k, v, o, cu_seqlens, max_seqlen): if _flash_attn_v3_available and (is_hopper() or is_4090()): try: flash_attention_v3_fwd(q, k, v, o, cu_seqlens, max_seqlen) + return except Exception as e: print(f"Failed to use sgl_kernel: {e}") _flash_attn_v3_available = False - else: - _flash_attention_triton_fwd(q, k, v, o, cu_seqlens, max_seqlen) + _flash_attention_triton_fwd(q, k, v, o, cu_seqlens, max_seqlen) diff --git a/unit_tests/models/vit/test_flash_attention_forward.py b/unit_tests/models/vit/test_flash_attention_forward.py index 70849445b4..f736ec4cd6 100644 --- a/unit_tests/models/vit/test_flash_attention_forward.py +++ b/unit_tests/models/vit/test_flash_attention_forward.py @@ -15,6 +15,58 @@ ) +def test_fa3_failure_computes_fallback_on_same_invocation(monkeypatch): + calls = [] + q = torch.zeros(2, 1, 4) + out = torch.full_like(q, -123) + cu = torch.tensor([0, 2], dtype=torch.int32) + + def fail_fa3(q, k, v, o, cu_seqlens, max_seqlen): + calls.append("fa3") + o[0].fill_(99) # A backend may fail after partially writing output. + raise RuntimeError("shape-dependent FA3 failure") + + def fallback(q, k, v, o, cu_seqlens, max_seqlen): + calls.append("triton") + o.fill_(7) + + monkeypatch.setattr(flashattention_nopad, "_flash_attn_v3_available", True) + monkeypatch.setattr(flashattention_nopad, "is_hopper", lambda: True) + monkeypatch.setattr(flashattention_nopad, "flash_attention_v3_fwd", fail_fa3, raising=False) + monkeypatch.setattr(flashattention_nopad, "_flash_attention_triton_fwd", fallback) + flash_attention_fwd(q, q, q, out, cu, 2) + assert calls == ["fa3", "triton"] + assert torch.equal(out, torch.full_like(q, 7)) + assert not flashattention_nopad._flash_attn_v3_available + + # Later calls should use the disabled-backend path without retrying FA3. + out.fill_(-123) + flash_attention_fwd(q, q, q, out, cu, 2) + assert calls == ["fa3", "triton", "triton"] + assert torch.equal(out, torch.full_like(q, 7)) + + +def test_fa3_success_does_not_run_fallback(monkeypatch): + q = torch.zeros(2, 1, 4) + out = torch.full_like(q, -123) + cu = torch.tensor([0, 2], dtype=torch.int32) + + def fa3(q, k, v, o, cu_seqlens, max_seqlen): + o.fill_(9) + + monkeypatch.setattr(flashattention_nopad, "_flash_attn_v3_available", True) + monkeypatch.setattr(flashattention_nopad, "is_hopper", lambda: True) + monkeypatch.setattr(flashattention_nopad, "flash_attention_v3_fwd", fa3, raising=False) + monkeypatch.setattr( + flashattention_nopad, + "_flash_attention_triton_fwd", + lambda *args: pytest.fail("a successful FA3 call must not run fallback"), + ) + flash_attention_fwd(q, q, q, out, cu, 2) + assert torch.equal(out, torch.full_like(q, 9)) + assert flashattention_nopad._flash_attn_v3_available + + def reference_attention_varlen(q, k, v, cu): """ q, k, v : (total_len, n_head, D) From 0bb017d8a533819e6e54faa22a7d61502f08572f Mon Sep 17 00:00:00 2001 From: sufubao Date: Sat, 17 Oct 2026 17:32:13 +0800 Subject: [PATCH 22/40] refactor: clarify recurrent execution contracts and freeze graph metadata --- .../common/basemodel/attention/linear/gdn.py | 37 +++--- .../common/basemodel/attention/linear/kda.py | 19 ++-- lightllm/common/basemodel/batch_objs.py | 2 +- lightllm/common/basemodel/infer_struct.py | 2 +- .../qwen3next_mem_manager.py | 2 + lightllm/common/recurrent_state.py | 57 ++++++---- lightllm/common/recurrent_state_types.py | 100 ++++++++++++++++ .../common/state_cache_manager/glm5_next.py | 1 + .../common/state_cache_manager/linear_att.py | 10 ++ .../linear_att/test_ssm_autotune.py | 10 +- .../common/test_recurrent_state_contract.py | 107 ++++++++++++++++++ unit_tests/models/glm5_next/test_mtp.py | 2 +- 12 files changed, 291 insertions(+), 58 deletions(-) create mode 100644 lightllm/common/recurrent_state_types.py create mode 100644 unit_tests/common/test_recurrent_state_contract.py diff --git a/lightllm/common/basemodel/attention/linear/gdn.py b/lightllm/common/basemodel/attention/linear/gdn.py index 571f85d8ac..64abcd0b75 100644 --- a/lightllm/common/basemodel/attention/linear/gdn.py +++ b/lightllm/common/basemodel/attention/linear/gdn.py @@ -1,5 +1,6 @@ import dataclasses import torch +from lightllm.common.recurrent_state_types import RecurrentStateContext from abc import ABC, abstractmethod from typing import TYPE_CHECKING from ..base_att import BaseAttBackend, BasePrefillAttState, BaseDecodeAttState, AttControl @@ -257,7 +258,7 @@ 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 self.backend.req_manager.recurrent_state.updates is not None: + if self.backend.req_manager.recurrent_state.update_backend is not None: self.b_ssm_buffer_idx = self.b_conv_buffer_idx return att_batch_size = self.b_conv_buffer_idx.shape[0] @@ -347,18 +348,17 @@ def _gdn_decode_kernel( backend.tp_num_v_heads, backend.head_v_dim, ) - recurrent_state = backend.req_manager.recurrent_state - if recurrent_state.updates is not None: - layer = ( - layer_weight.layer_num_ - - layer_weight.layer_num_ // backend.req_manager.linear_config.full_attention_interval - ) + updater = backend.req_manager.recurrent_state.update_backend + if updater is not None: + layer = backend.req_manager.linear_config.get_linear_layer_index(layer_weight.layer_num_) args = (query, key, value, a, b, layer_weight.linear_A_log.weight, layer_weight.linear_dt_bias.weight) if layer == 0: - self.infer_state.recurrent_state_context = recurrent_state.decode_context(*args) + self.infer_state.recurrent_state_context = RecurrentStateContext.from_run_config( + updater.get_run_config(*args) + ) return ( - recurrent_state.forward_updates( - layer, *args, self.b_conv_buffer_idx, context=self.infer_state.recurrent_state_context + updater.forward( + layer, *args, self.b_conv_buffer_idx, run_config=self.infer_state.recurrent_state_context.config ), z, ) @@ -407,21 +407,20 @@ def _gdn_mtp_kernel( ) query, key, value = backend._rearrange_mixed_qkv(mixed_qkv, decode=False) - recurrent_state = backend.req_manager.recurrent_state - if recurrent_state.updates is not None: - layer = ( - layer_weight.layer_num_ - - layer_weight.layer_num_ // backend.req_manager.linear_config.full_attention_interval - ) + updater = backend.req_manager.recurrent_state.update_backend + if updater is not None: + layer = backend.req_manager.linear_config.get_linear_layer_index(layer_weight.layer_num_) args = (query, key, value, a, b, layer_weight.linear_A_log.weight, layer_weight.linear_dt_bias.weight) if layer == 0: - self.infer_state.recurrent_state_context = recurrent_state.decode_context(*args, cu_seqlens_q) - return recurrent_state.forward_updates( + self.infer_state.recurrent_state_context = RecurrentStateContext.from_run_config( + updater.get_run_config(*args, cu_seqlens_q) + ) + return updater.forward( layer, *args, self.b_conv_buffer_idx, cu_seqlens_q, - context=self.infer_state.recurrent_state_context, + run_config=self.infer_state.recurrent_state_context.config, ) 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, diff --git a/lightllm/common/basemodel/attention/linear/kda.py b/lightllm/common/basemodel/attention/linear/kda.py index a896219d6a..e46fc37b69 100644 --- a/lightllm/common/basemodel/attention/linear/kda.py +++ b/lightllm/common/basemodel/attention/linear/kda.py @@ -8,6 +8,7 @@ from typing import TYPE_CHECKING import torch +from lightllm.common.recurrent_state_types import RecurrentStateContext from lightllm.common.basemodel.attention.base_att import ( AttControl, @@ -45,8 +46,6 @@ def __init__(self, model: "TpPartBaseModel"): self.conv_kernel_size = config["short_conv_kernel_size"] self.lower_bound = config.get("gate_lower_bound", -5.0) self.mtp_step = model.args.mtp_step - if model.req_manager.recurrent_state.updates is not None: - model.req_manager.recurrent_state.updates.lower_bound = self.lower_bound def create_att_prefill_state(self, infer_state: "InferStateInfo"): return KDAPrefillAttState(backend=self, infer_state=infer_state) @@ -184,7 +183,7 @@ def _init_fixed_mtp_decode_state(self, mtp_step: int): self._init_mtp_ssm_buffer_idx(mtp_size) def _init_mtp_ssm_buffer_idx(self, mtp_size: int): - if self.infer_state.req_manager.recurrent_state.updates is not None: + if self.infer_state.req_manager.recurrent_state.update_backend is not None: self.b_ssm_buffer_idx = self.b_conv_buffer_idx return att_batch_size = self.b_conv_buffer_idx.shape[0] @@ -232,9 +231,9 @@ def decode_att( v = v.view(*shape, backend.tp_num_heads, backend.head_dim) raw_gate = raw_gate.view(*shape, backend.tp_hidden_size) raw_beta = raw_beta.view(*shape, backend.tp_num_heads) - recurrent_state = self.infer_state.req_manager.recurrent_state - if recurrent_state.updates is not None: - layer = layer_num - layer_num // self.infer_state.req_manager.linear_config.full_attention_interval + updater = self.infer_state.req_manager.recurrent_state.update_backend + if updater is not None: + layer = self.infer_state.req_manager.linear_config.get_linear_layer_index(layer_num) args = ( q, k, @@ -245,15 +244,15 @@ def decode_att( layer_weight.linear_dt_bias.weight, ) if layer == 0: - self.infer_state.recurrent_state_context = recurrent_state.decode_context( - *args, self.b1_mtp_cu_q_seq_len + self.infer_state.recurrent_state_context = RecurrentStateContext.from_run_config( + updater.get_run_config(*args, self.b1_mtp_cu_q_seq_len) ) - return recurrent_state.forward_updates( + return updater.forward( layer, *args, self.b_conv_buffer_idx, self.b1_mtp_cu_q_seq_len, - context=self.infer_state.recurrent_state_context, + run_config=self.infer_state.recurrent_state_context.config, ) output, _ = fused_recurrent_kda( q=q, diff --git a/lightllm/common/basemodel/batch_objs.py b/lightllm/common/basemodel/batch_objs.py index 52218c143a..72a66a3229 100644 --- a/lightllm/common/basemodel/batch_objs.py +++ b/lightllm/common/basemodel/batch_objs.py @@ -7,7 +7,7 @@ from lightllm.utils.tensor_utils import tensor_to_no_ref_tensor if TYPE_CHECKING: - from lightllm.common.recurrent_state import RecurrentStateContext + from lightllm.common.recurrent_state_types import RecurrentStateContext @dataclass diff --git a/lightllm/common/basemodel/infer_struct.py b/lightllm/common/basemodel/infer_struct.py index 4fada1e7d8..f83825c1fb 100755 --- a/lightllm/common/basemodel/infer_struct.py +++ b/lightllm/common/basemodel/infer_struct.py @@ -14,7 +14,7 @@ from .attention import BasePrefillAttState, BaseDecodeAttState if TYPE_CHECKING: - from lightllm.common.recurrent_state import RecurrentStateContext + from lightllm.common.recurrent_state_types import RecurrentStateContext class InferStateInfo: 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 f9dae5b25f..e5365ea449 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,8 @@ def _free_linear_att_buffers(self): return def write_to_shm(self, req_manager): + # RecurrentState contains only IPC-safe state/operator configuration; + # worker execution resources must remain outside this shared object. self.recurrent_state = req_manager.recurrent_state # super().write_to_shm() 会用 ForkingPickler 序列化本对象，torch 在 dump 时会把 # CPU tensor 的 storage 原地迁到共享内存，使本进程大页 state cache 原本 diff --git a/lightllm/common/recurrent_state.py b/lightllm/common/recurrent_state.py index e0bb208efa..323c3e2b55 100644 --- a/lightllm/common/recurrent_state.py +++ b/lightllm/common/recurrent_state.py @@ -5,23 +5,16 @@ Attention backends retain responsibility for recurrent computation. """ -from dataclasses import dataclass -from typing import NamedTuple, Optional +from typing import TYPE_CHECKING, NamedTuple, Optional, Sequence, Tuple, Union import torch from lightllm.common.state_cache_manager import LayerCache, LinearAttCacheConfig +from lightllm.common.recurrent_state_types import RecurrentStateContext, SSMUpdateBackend - -@dataclass(frozen=True) -class RecurrentStateContext: - """Commit information for one forward/graph, passed unchanged by the runner. - - Compact reconstruction must use the verify kernel's reduction layout. - Keeping it per forward also permits overlapping graphs with different layouts. - """ - - _run_config: Optional[dict] = None +if TYPE_CHECKING: + from lightllm.common.basemodel.triton_kernel.linear_att.replayssm import ReplaySSMCache as _ReplaySSMCache + from lightllm.common.basemodel.triton_kernel.linear_att.replayssm_compact import CompactSSMCache as _CompactSSMCache class RecurrentStateSnapshot(NamedTuple): @@ -30,6 +23,16 @@ class RecurrentStateSnapshot(NamedTuple): class RecurrentState: + """Own request state and its lifecycle, without dispatching attention. + + PD currently shares this object through CUDA IPC as part of the memory + manager. Keep its fields limited to tensors, cache operators, and pickleable + static configuration. Graphs, streams, callbacks, batch contexts, and pinned + CPU checkpoint pools belong to their existing worker/runner owners. Adding + such resources requires an explicit transport view rather than serializing + them implicitly with this object. + """ + def __init__(self, max_request_num, linear_config: LinearAttCacheConfig, args, *, kda=False, device="cuda"): self.linear_config = linear_config self.mtp_step = args.mtp_step @@ -51,7 +54,7 @@ def __init__(self, max_request_num, linear_config: LinearAttCacheConfig, args, * layer_num=linear_config.linear_layer_num, device=device, ) - self.updates = None + self.updates: Optional[Union["_ReplaySSMCache", "_CompactSSMCache"]] = None if args.ssm_state_mode == "replay": from lightllm.common.basemodel.triton_kernel.linear_att.replayssm import ReplaySSMCache @@ -62,6 +65,7 @@ def __init__(self, max_request_num, linear_config: LinearAttCacheConfig, args, * linear_config.conv_state_dtype, num_key_heads=linear_config.num_linear_k_heads, kda=kda, + lower_bound=linear_config.gate_lower_bound, projection_mode=args.ssm_replay_config.get("compute_mode", "inline"), fold_method=args.ssm_replay_config.get("fold_method", "raw"), ) @@ -74,8 +78,14 @@ def __init__(self, max_request_num, linear_config: LinearAttCacheConfig, args, * linear_config.conv_state_dtype, num_key_heads=linear_config.num_linear_k_heads, kda=kda, + lower_bound=linear_config.gate_lower_bound, ) + @property + def update_backend(self) -> Optional[SSMUpdateBackend]: + """Optional compact/replay execution backend for attention to call.""" + return self.updates + def _clear_positions(self, req_idx): if self.updates is not None: self.updates.clear_history(req_idx) @@ -102,13 +112,15 @@ def prepare_decode(self, reqs, cu_seqlens=None, *, single_token=False): # guaranteed token, including when forward is graph-replayed. self.updates.accept_updates(reqs) - def decode_context(self, *args): - return RecurrentStateContext(self.updates.get_run_config(*args)) - - def forward_updates(self, layer, *args, context): - return self.updates.forward(layer, *args, run_config=context._run_config) - - def commit(self, b_req_mtp_start_loc, b_req_idx, b_mtp_index, accepted_index, verify_width, batches=None): + def commit( + self, + b_req_mtp_start_loc, + b_req_idx, + b_mtp_index, + accepted_index, + verify_width, + batches: Optional[Sequence[Tuple[int, Optional[RecurrentStateContext]]]] = None, + ): """Commit accepted prefixes without requiring full-state materialization. ``batches`` partitions logical requests in the same order as the MTP @@ -143,7 +155,7 @@ def commit(self, b_req_mtp_start_loc, b_req_idx, b_mtp_index, accepted_index, ve self.updates.accept_updates( reqs[start : start + size], self.accepted_index, - run_config=context._run_config if context is not None else None, + run_config=context.config if context is not None else None, ) start += size @@ -218,6 +230,5 @@ def export(self, req_idx, conv, ssm): ssm.copy_(snapshot.ssm, non_blocking=True) def layer_buffers(self, layer_idx_in_all): - assert 0 <= layer_idx_in_all < self.linear_config.all_layer_num - layer = layer_idx_in_all - layer_idx_in_all // self.linear_config.full_attention_interval + layer = self.linear_config.get_linear_layer_index(layer_idx_in_all) return self.conv_state.buffer[layer], self.ssm_state.buffer[layer] diff --git a/lightllm/common/recurrent_state_types.py b/lightllm/common/recurrent_state_types.py new file mode 100644 index 0000000000..69e25227e1 --- /dev/null +++ b/lightllm/common/recurrent_state_types.py @@ -0,0 +1,100 @@ +"""Immutable execution metadata and the attention-facing SSM updater contract.""" + +from collections.abc import Iterator, Mapping +from dataclasses import dataclass +from typing import Optional, Protocol + +import torch + + +@dataclass(frozen=True, slots=True) +class SSMRunConfig(Mapping[str, int]): + """A detached kernel layout, also readable by existing Triton launch code.""" + + BV: int + num_warps: int + num_stages: int = 3 + + def __post_init__(self): + for key in self: + if type(self[key]) is not int: + raise TypeError(f"SSM kernel configuration {key} must be an integer") + + @classmethod + def from_mapping(cls, config: Mapping[str, int]) -> "SSMRunConfig": + unknown = config.keys() - {"BV", "num_warps", "num_stages"} + if unknown: + raise ValueError(f"Unknown SSM kernel configuration fields: {unknown}") + return cls(BV=config["BV"], num_warps=config["num_warps"], num_stages=config.get("num_stages", 3)) + + def __getitem__(self, key: str) -> int: + if key == "BV": + return self.BV + if key == "num_warps": + return self.num_warps + if key == "num_stages": + return self.num_stages + raise KeyError(key) + + def __iter__(self) -> Iterator[str]: + return iter(("BV", "num_warps", "num_stages")) + + def __len__(self) -> int: + return 3 + + +@dataclass(frozen=True, slots=True) +class RecurrentStateContext: + """Per-forward/graph commit metadata, passed unchanged by the runner. + + Both verify and compact reconstruction read the same immutable layout. + Publishing a context detaches it from the tuner's mutable dictionaries. + """ + + config: SSMRunConfig + + def __post_init__(self): + if not isinstance(self.config, SSMRunConfig): + raise TypeError("Use from_run_config() to freeze a mutable kernel configuration") + + @classmethod + def from_run_config(cls, config: Mapping[str, int]) -> "RecurrentStateContext": + return cls(SSMRunConfig.from_mapping(config)) + + +class SSMUpdateBackend(Protocol): + """Attention executes the updater directly; lifecycle stays with its owner. + + ``a`` and ``b`` are raw gate/beta projections. GDN uses scalar decay per + value head; KDA's ``a`` and ``bias`` include the key dimension. A missing + ``cu_seqlens`` denotes ordinary one-token-per-request decode. + """ + + def get_run_config( + self, + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + a: torch.Tensor, + b: torch.Tensor, + a_log: torch.Tensor, + bias: torch.Tensor, + cu_seqlens: Optional[torch.Tensor] = None, + ) -> Mapping[str, int]: + ... + + def forward( + self, + layer: int, + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + a: torch.Tensor, + b: torch.Tensor, + a_log: torch.Tensor, + bias: torch.Tensor, + reqs: torch.Tensor, + cu_seqlens: Optional[torch.Tensor] = None, + run_config: Optional[Mapping[str, int]] = None, + ) -> torch.Tensor: + ... diff --git a/lightllm/common/state_cache_manager/glm5_next.py b/lightllm/common/state_cache_manager/glm5_next.py index 13195a74ac..77c42905c1 100644 --- a/lightllm/common/state_cache_manager/glm5_next.py +++ b/lightllm/common/state_cache_manager/glm5_next.py @@ -47,6 +47,7 @@ def from_model_config(cls, config, args): full_attention_interval=4, all_layer_num=layers, draft_full_att_kv_layer_num=get_added_mtp_kv_layer_num() if args.mtp_mode is not None else 0, + gate_lower_bound=linear.get("gate_lower_bound", -5.0), index_kpool=config["index_kpool"], index_head_dim=config["index_head_dim"], ) diff --git a/lightllm/common/state_cache_manager/linear_att.py b/lightllm/common/state_cache_manager/linear_att.py index 35a37dcd28..e533537deb 100644 --- a/lightllm/common/state_cache_manager/linear_att.py +++ b/lightllm/common/state_cache_manager/linear_att.py @@ -36,6 +36,8 @@ class LinearAttCacheConfig: full_attention_interval: int all_layer_num: int # 包括 linear att 和 full att 的层加起来的层数 draft_full_att_kv_layer_num: int = 0 + # Static KDA recurrence parameter; unused by GDN. + gate_lower_bound: float = -5.0 def get_conv_dim(self): # 第一项对应q的参数，第二项对应k的参数，第三项对应v的参数 @@ -69,6 +71,14 @@ def get_full_att_kv_layer_index(self, layer_index: int) -> int: ) return kv_layer_index + def get_linear_layer_index(self, layer_index: int) -> int: + """Map a target linear-attention layer to its packed recurrent-state row.""" + assert 0 <= layer_index < self.all_layer_num, f"invalid transformer layer index {layer_index}" + assert (layer_index + 1) % self.full_attention_interval != 0, f"layer {layer_index} uses full attention" + linear_index = layer_index - layer_index // self.full_attention_interval + assert linear_index < self.linear_layer_num, f"layer {layer_index} maps outside the recurrent-state pool" + return linear_index + def get_conv_state_shape(self): # Base committed sliding-window state, without speculative MTP tail. return (self.get_conv_dim(), self.conv_kernel_size - 1) diff --git a/unit_tests/common/basemodel/triton_kernel/linear_att/test_ssm_autotune.py b/unit_tests/common/basemodel/triton_kernel/linear_att/test_ssm_autotune.py index b8e6329d68..cbf41fa38f 100644 --- a/unit_tests/common/basemodel/triton_kernel/linear_att/test_ssm_autotune.py +++ b/unit_tests/common/basemodel/triton_kernel/linear_att/test_ssm_autotune.py @@ -337,7 +337,7 @@ def test_two_verify_graphs_accept_with_their_own_configs(monkeypatch, mode, proj assert config == configs[index] ref = make_inputs(mode, dtype, "cuda", projection_mode=projection, kda=kda)["cache"] ref.state.copy_(initial) - ref.run_config = config + ref.run_config = dict(config) ref._config_is_fixed = True def forward(target): @@ -351,15 +351,19 @@ def forward(target): state = SimpleNamespace( hidden_collector=SimpleNamespace(finish_output=lambda **kw: ModelMtpOutputCollector()), prompt_logics=None, - recurrent_state_context=RecurrentStateContext(config), + recurrent_state_context=RecurrentStateContext.from_run_config(config), logits_row_stride=1, ) output = TpPartBaseModel._create_model_output( None, PostLayerOutput(logits=outputs[0].view(count * 4, -1)), state ) output = TpPartBaseModel._create_unpad_decode_model_output(None, output, len(ids) * 4) - assert output.recurrent_state_context._run_config is config + assert dict(output.recurrent_state_context.config) == config buckets.append((graph, outputs, output, ref, args, reqs, cu, ids)) + # Published contexts and captured launches must survive later changes to + # the tuner's dictionaries, including when two graphs replay alternately. + for config in configs: + config.update(BV=128, num_warps=4, num_stages=3) cache.state.copy_(initial) if mode == "replay": cache.cursors.zero_() diff --git a/unit_tests/common/test_recurrent_state_contract.py b/unit_tests/common/test_recurrent_state_contract.py new file mode 100644 index 0000000000..82b701f0e7 --- /dev/null +++ b/unit_tests/common/test_recurrent_state_contract.py @@ -0,0 +1,107 @@ +import pickle +from dataclasses import FrozenInstanceError +from types import SimpleNamespace + +import pytest +import torch + +from lightllm.common.recurrent_state import RecurrentState +from lightllm.common.recurrent_state_types import RecurrentStateContext +from lightllm.common.state_cache_manager import Glm5NextCacheConfig, LinearAttCacheConfig +from lightllm.server.core.objs.start_args_type import StartArgs + + +def test_published_context_is_detached_immutable_and_pickleable(): + tuning = {"BV": 8, "num_warps": 1} + context = RecurrentStateContext.from_run_config(tuning) + tuning.update(BV=128, num_warps=4, num_stages=1) + assert dict(context.config) == {"BV": 8, "num_warps": 1, "num_stages": 3} + with pytest.raises(FrozenInstanceError): + context.config.BV = 128 + with pytest.raises(TypeError): + context.config["BV"] = 128 + with pytest.raises(TypeError, match="from_run_config"): + RecurrentStateContext(tuning) + restored = pickle.loads(pickle.dumps(context)) + assert restored == context + assert dict(restored.config) == dict(context.config) + + +def test_context_rejects_unrepresented_kernel_options(): + with pytest.raises(ValueError, match="Unknown SSM kernel configuration fields"): + RecurrentStateContext.from_run_config({"BV": 32, "num_warps": 1, "future_layout_option": 2}) + with pytest.raises(TypeError, match="must be an integer"): + RecurrentStateContext.from_run_config({"BV": [32], "num_warps": 1}) + + +@pytest.fixture +def linear_config(): + return LinearAttCacheConfig( + tp_world_size=1, + full_att_all_num_kv_heads=1, + full_att_dtype=torch.bfloat16, + full_att_num_kv_heads=1, + full_att_head_dim=32, + global_linear_k_heads=1, + global_linear_v_heads=2, + num_linear_k_heads=1, + num_linear_v_heads=2, + head_linear_k_dim=32, + head_linear_v_dim=32, + conv_kernel_size=4, + linear_layer_num=6, + conv_state_dtype=torch.bfloat16, + ssm_state_dtype=torch.float32, + full_attention_interval=3, + all_layer_num=8, + ) + + +@pytest.mark.parametrize("layer,row", [(0, 0), (1, 1), (3, 2), (4, 3), (6, 4), (7, 5)]) +def test_layer_buffers_follow_packed_linear_layers(linear_config, layer, row): + owner = object.__new__(RecurrentState) + owner.linear_config = linear_config + conv = torch.arange(6).view(6, 1) + ssm = conv + 10 + owner.conv_state = SimpleNamespace(buffer=conv) + owner.ssm_state = SimpleNamespace(buffer=ssm) + assert linear_config.get_linear_layer_index(layer) == row + actual_conv, actual_ssm = owner.layer_buffers(layer) + assert actual_conv.item() == row + assert actual_ssm.item() == row + 10 + + +@pytest.mark.parametrize("layer", [-1, 2, 5, 8]) +def test_linear_mapping_rejects_full_attention_and_out_of_range_layers(linear_config, layer): + with pytest.raises(AssertionError): + linear_config.get_linear_layer_index(layer) + + +def test_model_configuration_retains_nondefault_kda_gate_bound(): + config = Glm5NextCacheConfig.from_model_config( + { + "kv_lora_rank": 512, + "num_hidden_layers": 4, + "layer_types": ["linear_attention"] * 3 + ["deepseek_sparse_attention"], + "index_kpool": 4, + "index_head_dim": 128, + "linear_attn_config": { + "num_heads": 8, + "head_dim": 128, + "short_conv_kernel_size": 4, + "kda_layers": [0, 1, 2], + "gate_lower_bound": -2.5, + }, + }, + StartArgs(tp=4, data_type="bfloat16"), + ) + assert config.gate_lower_bound == -2.5 + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") +@pytest.mark.parametrize("mode", ["compact", "replay"]) +def test_kda_updater_receives_gate_bound_at_construction(linear_config, mode): + linear_config.gate_lower_bound = -2.5 + args = SimpleNamespace(mtp_step=2, ssm_state_mode=mode, ssm_replay_config={}) + owner = RecurrentState(2, linear_config, args, kda=True) + assert owner.updates.lower_bound == -2.5 diff --git a/unit_tests/models/glm5_next/test_mtp.py b/unit_tests/models/glm5_next/test_mtp.py index 1bd3145077..82ab899161 100644 --- a/unit_tests/models/glm5_next/test_mtp.py +++ b/unit_tests/models/glm5_next/test_mtp.py @@ -199,7 +199,7 @@ def test_kda_backend_routes_accepted_conv_and_ssm_states(dynamic, mode, mtp_step recurrent_state=recurrent_state, HOLD_REQUEST_ID=3, get_mamba_cache=lambda layer: (conv, ssm), - linear_config=SimpleNamespace(full_attention_interval=4), + linear_config=SimpleNamespace(get_linear_layer_index=lambda layer: 0), ) if mode != "native": from lightllm.common.basemodel.triton_kernel.linear_att.replayssm import ReplaySSMCache From 600df9ee4e81f06a90cad754df19ed47d647e762 Mon Sep 17 00:00:00 2001 From: sufubao Date: Mon, 19 Oct 2026 15:44:34 +0800 Subject: [PATCH 23/40] refactor: split independent vision FA3 fixes out of ReplaySSM --- .../common/basemodel/attention_vit/fa3/fp.py | 2 - .../vit/triton_kernel/flashattention_nopad.py | 6 +- .../vit/test_flash_attention_forward.py | 82 +------------------ 3 files changed, 5 insertions(+), 85 deletions(-) diff --git a/lightllm/common/basemodel/attention_vit/fa3/fp.py b/lightllm/common/basemodel/attention_vit/fa3/fp.py index a8e56417c5..f804116f1f 100644 --- a/lightllm/common/basemodel/attention_vit/fa3/fp.py +++ b/lightllm/common/basemodel/attention_vit/fa3/fp.py @@ -53,8 +53,6 @@ def _vit_att_fwd( pack_gqa=None, sm_margin=0, sinks=None, - sparse_mask_fine=None, - only_qv=False, ) return o diff --git a/lightllm/models/vit/triton_kernel/flashattention_nopad.py b/lightllm/models/vit/triton_kernel/flashattention_nopad.py index af062c2a5b..b21464a02b 100644 --- a/lightllm/models/vit/triton_kernel/flashattention_nopad.py +++ b/lightllm/models/vit/triton_kernel/flashattention_nopad.py @@ -203,8 +203,6 @@ def flash_attention_v3_fwd( pack_gqa=None, sm_margin=0, sinks=None, - sparse_mask_fine=None, - only_qv=False, ) return o @@ -222,8 +220,8 @@ def flash_attention_fwd(q, k, v, o, cu_seqlens, max_seqlen): if _flash_attn_v3_available and (is_hopper() or is_4090()): try: flash_attention_v3_fwd(q, k, v, o, cu_seqlens, max_seqlen) - return except Exception as e: print(f"Failed to use sgl_kernel: {e}") _flash_attn_v3_available = False - _flash_attention_triton_fwd(q, k, v, o, cu_seqlens, max_seqlen) + else: + _flash_attention_triton_fwd(q, k, v, o, cu_seqlens, max_seqlen) diff --git a/unit_tests/models/vit/test_flash_attention_forward.py b/unit_tests/models/vit/test_flash_attention_forward.py index f736ec4cd6..53b70e6836 100644 --- a/unit_tests/models/vit/test_flash_attention_forward.py +++ b/unit_tests/models/vit/test_flash_attention_forward.py @@ -3,68 +3,6 @@ import time import pytest from lightllm.models.vit.triton_kernel.flashattention_nopad import flash_attention_fwd -from lightllm.models.vit.triton_kernel import flashattention_nopad -from lightllm.common.basemodel.attention_vit.fa3.fp import Fa3VitAttBackend - - -_fa3_mark = pytest.mark.skipif( - not flashattention_nopad._flash_attn_v3_available - or not torch.cuda.is_available() - or torch.cuda.get_device_capability()[0] < 9, - reason="FA3 requires sgl_kernel and a Hopper GPU", -) - - -def test_fa3_failure_computes_fallback_on_same_invocation(monkeypatch): - calls = [] - q = torch.zeros(2, 1, 4) - out = torch.full_like(q, -123) - cu = torch.tensor([0, 2], dtype=torch.int32) - - def fail_fa3(q, k, v, o, cu_seqlens, max_seqlen): - calls.append("fa3") - o[0].fill_(99) # A backend may fail after partially writing output. - raise RuntimeError("shape-dependent FA3 failure") - - def fallback(q, k, v, o, cu_seqlens, max_seqlen): - calls.append("triton") - o.fill_(7) - - monkeypatch.setattr(flashattention_nopad, "_flash_attn_v3_available", True) - monkeypatch.setattr(flashattention_nopad, "is_hopper", lambda: True) - monkeypatch.setattr(flashattention_nopad, "flash_attention_v3_fwd", fail_fa3, raising=False) - monkeypatch.setattr(flashattention_nopad, "_flash_attention_triton_fwd", fallback) - flash_attention_fwd(q, q, q, out, cu, 2) - assert calls == ["fa3", "triton"] - assert torch.equal(out, torch.full_like(q, 7)) - assert not flashattention_nopad._flash_attn_v3_available - - # Later calls should use the disabled-backend path without retrying FA3. - out.fill_(-123) - flash_attention_fwd(q, q, q, out, cu, 2) - assert calls == ["fa3", "triton", "triton"] - assert torch.equal(out, torch.full_like(q, 7)) - - -def test_fa3_success_does_not_run_fallback(monkeypatch): - q = torch.zeros(2, 1, 4) - out = torch.full_like(q, -123) - cu = torch.tensor([0, 2], dtype=torch.int32) - - def fa3(q, k, v, o, cu_seqlens, max_seqlen): - o.fill_(9) - - monkeypatch.setattr(flashattention_nopad, "_flash_attn_v3_available", True) - monkeypatch.setattr(flashattention_nopad, "is_hopper", lambda: True) - monkeypatch.setattr(flashattention_nopad, "flash_attention_v3_fwd", fa3, raising=False) - monkeypatch.setattr( - flashattention_nopad, - "_flash_attention_triton_fwd", - lambda *args: pytest.fail("a successful FA3 call must not run fallback"), - ) - flash_attention_fwd(q, q, q, out, cu, 2) - assert torch.equal(out, torch.full_like(q, 9)) - assert flashattention_nopad._flash_attn_v3_available def reference_attention_varlen(q, k, v, cu): @@ -94,21 +32,7 @@ def reference_attention_varlen(q, k, v, cu): @pytest.mark.parametrize("dtype,atol", [(torch.float16, 1e-2), (torch.bfloat16, 2e-2)]) -@pytest.mark.parametrize( - "attention_fn", - [ - flash_attention_fwd, - pytest.param( - getattr(flashattention_nopad, "flash_attention_v3_fwd", None), - marks=_fa3_mark, - ), - pytest.param( - Fa3VitAttBackend._vit_att_fwd, - marks=_fa3_mark, - ), - ], -) -def test_varlen(dtype, atol, attention_fn, batch=4, heads=8, d=80, device="cuda:0"): +def test_varlen(dtype, atol, batch=4, heads=8, d=80, device="cuda:0"): torch.manual_seed(0) lengths = torch.randint(1, 257, (batch,)) max_len = int(lengths.max().item()) @@ -121,10 +45,10 @@ def test_varlen(dtype, atol, attention_fn, batch=4, heads=8, d=80, device="cuda: k = torch.randn_like(q) v = torch.randn_like(q) out_tri = torch.randn_like(q) - attention_fn(q, k, v, out_tri, cu, max_len) + flash_attention_fwd(q, k, v, out_tri, cu, max_len) a = time.time() for _ in range(100): - attention_fn(q, k, v, out_tri, cu, max_len) + flash_attention_fwd(q, k, v, out_tri, cu, max_len) b = time.time() print(f"flash_attention_fwd time: {(b - a) / 100 * 1000:.2f} ms") out_ref = reference_attention_varlen(q, k, v, cu) From 60ba1f38cec574a170fd049086049f2c7282993c Mon Sep 17 00:00:00 2001 From: sufubao Date: Mon, 19 Oct 2026 16:47:02 +0800 Subject: [PATCH 24/40] fix: shorten GDN precompute tuning cache filenames --- .../basemodel/triton_kernel/linear_att/ssm_autotune.py | 2 +- .../triton_kernel/linear_att/test_ssm_autotune.py | 9 +++++---- 2 files changed, 6 insertions(+), 5 deletions(-) diff --git a/lightllm/common/basemodel/triton_kernel/linear_att/ssm_autotune.py b/lightllm/common/basemodel/triton_kernel/linear_att/ssm_autotune.py index 681935f953..cf623a7a1e 100644 --- a/lightllm/common/basemodel/triton_kernel/linear_att/ssm_autotune.py +++ b/lightllm/common/basemodel/triton_kernel/linear_att/ssm_autotune.py @@ -42,7 +42,7 @@ def static_key(cache, mode, q, k, v, a, b, cu_seqlens): key["fold"] = cache.fold_method return key key["projection"] = ( - ("kda_mma" if cache.kda else "block_mma_activation_precision") + ("kda_mma" if cache.kda else "gdn_mma") if cache.projection_mode == "precompute" and cache.verify_width > 1 else "inline" ) diff --git a/unit_tests/common/basemodel/triton_kernel/linear_att/test_ssm_autotune.py b/unit_tests/common/basemodel/triton_kernel/linear_att/test_ssm_autotune.py index cbf41fa38f..419cb47e81 100644 --- a/unit_tests/common/basemodel/triton_kernel/linear_att/test_ssm_autotune.py +++ b/unit_tests/common/basemodel/triton_kernel/linear_att/test_ssm_autotune.py @@ -246,10 +246,11 @@ def test_kda_rule_and_gate_bound_separate_tuning_cache(mode): @pytest.mark.parametrize("tp", [1, 2, 4, 8]) @pytest.mark.parametrize("dtype", [torch.float32, torch.bfloat16]) -def test_kda_production_tuning_cache_can_be_written(monkeypatch, tmp_path, tp, dtype): +@pytest.mark.parametrize("kda", [False, True], ids=["gdn", "kda"]) +def test_production_tuning_cache_can_be_written(monkeypatch, tmp_path, tp, dtype, kda): from types import SimpleNamespace - monkeypatch.setattr(kernel_config_module, "get_current_device_name", lambda: "NVIDIA H100 80GB HBM3") + monkeypatch.setattr(kernel_config_module, "get_current_device_name", lambda: "NVIDIA H100 80GB HBM3 MIG 3g.40gb") heads = 64 // tp cache = SimpleNamespace( state=SimpleNamespace(shape=(34, 1, heads, 128, 128), dtype=dtype), @@ -257,13 +258,13 @@ def test_kda_production_tuning_cache_can_be_written(monkeypatch, tmp_path, tp, d verify_width=3, num_key_heads=heads, capacity=8, - kda=True, + kda=kda, lower_bound=-5.0, projection_mode="precompute", ) packed = torch.empty(1, 3, heads * 128 * 3, dtype=torch.bfloat16) q, k, v = [x.view(1, 3, heads, 128) for x in packed.chunk(3, -1)] - a = torch.empty(3, heads * 128, dtype=q.dtype) + a = torch.empty(3, heads * 128 if kda else heads, dtype=q.dtype) b = torch.empty(3, heads, dtype=q.dtype) key = tuning.static_key(cache, "replay", q, k, v, a, b, torch.tensor([0, 3])) path = tmp_path / KernelConfigs.get_config_file_name(key) From a1c86f02d8d4ebd9980f8c17d1876359c01a7c0d Mon Sep 17 00:00:00 2001 From: sufubao Date: Mon, 19 Oct 2026 17:19:32 +0800 Subject: [PATCH 25/40] refactor: clarify linear layer index naming --- lightllm/common/basemodel/attention/linear/gdn.py | 15 +++++++++------ lightllm/common/basemodel/attention/linear/kda.py | 6 +++--- 2 files changed, 12 insertions(+), 9 deletions(-) diff --git a/lightllm/common/basemodel/attention/linear/gdn.py b/lightllm/common/basemodel/attention/linear/gdn.py index 3c3a6f6af5..4357e4ad60 100644 --- a/lightllm/common/basemodel/attention/linear/gdn.py +++ b/lightllm/common/basemodel/attention/linear/gdn.py @@ -351,15 +351,18 @@ def _gdn_decode_kernel( ) updater = backend.req_manager.recurrent_state.update_backend if updater is not None: - layer = backend.req_manager.linear_config.get_linear_layer_index(layer_weight.layer_num_) + linear_layer_idx = backend.req_manager.linear_config.get_linear_layer_index(layer_weight.layer_num_) args = (query, key, value, a, b, layer_weight.linear_A_log.weight, layer_weight.linear_dt_bias.weight) - if layer == 0: + if linear_layer_idx == 0: self.infer_state.recurrent_state_context = RecurrentStateContext.from_run_config( updater.get_run_config(*args) ) return ( updater.forward( - layer, *args, self.b_conv_buffer_idx, run_config=self.infer_state.recurrent_state_context.config + linear_layer_idx, + *args, + self.b_conv_buffer_idx, + run_config=self.infer_state.recurrent_state_context.config, ), z, ) @@ -410,14 +413,14 @@ def _gdn_mtp_kernel( query, key, value = backend._rearrange_mixed_qkv(mixed_qkv, decode=False) updater = backend.req_manager.recurrent_state.update_backend if updater is not None: - layer = backend.req_manager.linear_config.get_linear_layer_index(layer_weight.layer_num_) + linear_layer_idx = backend.req_manager.linear_config.get_linear_layer_index(layer_weight.layer_num_) args = (query, key, value, a, b, layer_weight.linear_A_log.weight, layer_weight.linear_dt_bias.weight) - if layer == 0: + if linear_layer_idx == 0: self.infer_state.recurrent_state_context = RecurrentStateContext.from_run_config( updater.get_run_config(*args, cu_seqlens_q) ) return updater.forward( - layer, + linear_layer_idx, *args, self.b_conv_buffer_idx, cu_seqlens_q, diff --git a/lightllm/common/basemodel/attention/linear/kda.py b/lightllm/common/basemodel/attention/linear/kda.py index eb50c1de68..dd27180cf2 100644 --- a/lightllm/common/basemodel/attention/linear/kda.py +++ b/lightllm/common/basemodel/attention/linear/kda.py @@ -234,7 +234,7 @@ def decode_att( raw_beta = raw_beta.view(*shape, backend.tp_num_heads) updater = self.infer_state.req_manager.recurrent_state.update_backend if updater is not None: - layer = self.infer_state.req_manager.linear_config.get_linear_layer_index(layer_num) + linear_layer_idx = self.infer_state.req_manager.linear_config.get_linear_layer_index(layer_num) args = ( q, k, @@ -244,12 +244,12 @@ def decode_att( layer_weight.linear_A_log.weight, layer_weight.linear_dt_bias.weight, ) - if layer == 0: + if linear_layer_idx == 0: self.infer_state.recurrent_state_context = RecurrentStateContext.from_run_config( updater.get_run_config(*args, self.b1_mtp_cu_q_seq_len) ) return updater.forward( - layer, + linear_layer_idx, *args, self.b_conv_buffer_idx, self.b1_mtp_cu_q_seq_len, From e1dc47acb52ec104e03b3b5a6a71f9b6c52c8940 Mon Sep 17 00:00:00 2001 From: sufubao Date: Mon, 19 Oct 2026 17:22:20 +0800 Subject: [PATCH 26/40] refactor: pass immutable SSM run configs directly --- .../common/basemodel/attention/linear/gdn.py | 14 ++++------ .../common/basemodel/attention/linear/kda.py | 6 ++--- lightllm/common/basemodel/basemodel.py | 2 +- lightllm/common/basemodel/batch_objs.py | 6 ++--- lightllm/common/basemodel/infer_struct.py | 4 +-- lightllm/common/recurrent_state.py | 10 +++---- lightllm/common/recurrent_state_types.py | 19 -------------- .../mode_backend/chunked_prefill/impl.py | 16 +++++++++--- .../mode_backend/dp_backend/impl.py | 6 ++--- .../common/basemodel/test_draft_tail_graph.py | 2 +- .../linear_att/test_ssm_autotune.py | 9 ++++--- .../common/test_recurrent_state_contract.py | 26 +++++++++---------- unit_tests/models/test_post_layer_output.py | 2 +- 13 files changed, 54 insertions(+), 68 deletions(-) diff --git a/lightllm/common/basemodel/attention/linear/gdn.py b/lightllm/common/basemodel/attention/linear/gdn.py index 4357e4ad60..283df27afe 100644 --- a/lightllm/common/basemodel/attention/linear/gdn.py +++ b/lightllm/common/basemodel/attention/linear/gdn.py @@ -1,7 +1,7 @@ import dataclasses import weakref import torch -from lightllm.common.recurrent_state_types import RecurrentStateContext +from lightllm.common.recurrent_state_types import SSMRunConfig from abc import ABC, abstractmethod from typing import TYPE_CHECKING from ..base_att import BaseAttBackend, BasePrefillAttState, BaseDecodeAttState, AttControl @@ -354,15 +354,13 @@ def _gdn_decode_kernel( linear_layer_idx = backend.req_manager.linear_config.get_linear_layer_index(layer_weight.layer_num_) args = (query, key, value, a, b, layer_weight.linear_A_log.weight, layer_weight.linear_dt_bias.weight) if linear_layer_idx == 0: - self.infer_state.recurrent_state_context = RecurrentStateContext.from_run_config( - updater.get_run_config(*args) - ) + self.infer_state.ssm_run_config = SSMRunConfig.from_mapping(updater.get_run_config(*args)) return ( updater.forward( linear_layer_idx, *args, self.b_conv_buffer_idx, - run_config=self.infer_state.recurrent_state_context.config, + run_config=self.infer_state.ssm_run_config, ), z, ) @@ -416,15 +414,13 @@ def _gdn_mtp_kernel( linear_layer_idx = backend.req_manager.linear_config.get_linear_layer_index(layer_weight.layer_num_) args = (query, key, value, a, b, layer_weight.linear_A_log.weight, layer_weight.linear_dt_bias.weight) if linear_layer_idx == 0: - self.infer_state.recurrent_state_context = RecurrentStateContext.from_run_config( - updater.get_run_config(*args, cu_seqlens_q) - ) + self.infer_state.ssm_run_config = SSMRunConfig.from_mapping(updater.get_run_config(*args, cu_seqlens_q)) return updater.forward( linear_layer_idx, *args, self.b_conv_buffer_idx, cu_seqlens_q, - run_config=self.infer_state.recurrent_state_context.config, + run_config=self.infer_state.ssm_run_config, ) 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, diff --git a/lightllm/common/basemodel/attention/linear/kda.py b/lightllm/common/basemodel/attention/linear/kda.py index dd27180cf2..a2fe2288c0 100644 --- a/lightllm/common/basemodel/attention/linear/kda.py +++ b/lightllm/common/basemodel/attention/linear/kda.py @@ -9,7 +9,7 @@ from typing import TYPE_CHECKING import torch -from lightllm.common.recurrent_state_types import RecurrentStateContext +from lightllm.common.recurrent_state_types import SSMRunConfig from lightllm.common.basemodel.attention.base_att import ( AttControl, @@ -245,7 +245,7 @@ def decode_att( layer_weight.linear_dt_bias.weight, ) if linear_layer_idx == 0: - self.infer_state.recurrent_state_context = RecurrentStateContext.from_run_config( + self.infer_state.ssm_run_config = SSMRunConfig.from_mapping( updater.get_run_config(*args, self.b1_mtp_cu_q_seq_len) ) return updater.forward( @@ -253,7 +253,7 @@ def decode_att( *args, self.b_conv_buffer_idx, self.b1_mtp_cu_q_seq_len, - run_config=self.infer_state.recurrent_state_context.config, + run_config=self.infer_state.ssm_run_config, ) output, _ = fused_recurrent_kda( q=q, diff --git a/lightllm/common/basemodel/basemodel.py b/lightllm/common/basemodel/basemodel.py index a62cfcb9c8..0d43e6d9dd 100755 --- a/lightllm/common/basemodel/basemodel.py +++ b/lightllm/common/basemodel/basemodel.py @@ -530,7 +530,7 @@ def _create_model_output(self, post_output: PostLayerOutput, infer_state: InferS mtp_collector=infer_state.hidden_collector.finish_output(infer_state=infer_state), prompt_logics=infer_state.prompt_logics, logits_token_ids=post_output.logits_token_ids, - recurrent_state_context=infer_state.recurrent_state_context, + ssm_run_config=infer_state.ssm_run_config, logits_row_stride=infer_state.logits_row_stride, ) return output diff --git a/lightllm/common/basemodel/batch_objs.py b/lightllm/common/basemodel/batch_objs.py index 72a66a3229..db678103e1 100644 --- a/lightllm/common/basemodel/batch_objs.py +++ b/lightllm/common/basemodel/batch_objs.py @@ -7,7 +7,7 @@ from lightllm.utils.tensor_utils import tensor_to_no_ref_tensor if TYPE_CHECKING: - from lightllm.common.recurrent_state_types import RecurrentStateContext + from lightllm.common.recurrent_state_types import SSMRunConfig @dataclass @@ -210,8 +210,8 @@ class ModelOutput: # target 模型会将候选值回填到完整词表 logits，因此不需要该映射。 logits_token_ids: Optional[torch.Tensor] = None - # Opaque commit context for this forward/graph; preserved when unpadding. - recurrent_state_context: Optional["RecurrentStateContext"] = None + # Immutable SSM configuration for this forward/graph; preserved when unpadding. + ssm_run_config: Optional["SSMRunConfig"] = None # Number of physical decode rows represented by each selected logits row. logits_row_stride: int = 1 diff --git a/lightllm/common/basemodel/infer_struct.py b/lightllm/common/basemodel/infer_struct.py index 5dcd7ab12d..0f94a34a10 100755 --- a/lightllm/common/basemodel/infer_struct.py +++ b/lightllm/common/basemodel/infer_struct.py @@ -14,7 +14,7 @@ from .attention import BasePrefillAttState, BaseDecodeAttState if TYPE_CHECKING: - from lightllm.common.recurrent_state_types import RecurrentStateContext + from lightllm.common.recurrent_state_types import SSMRunConfig class InferStateInfo: @@ -68,7 +68,7 @@ def __init__(self): # 在microbatch overlap的运行模式下，用于标记当前 microbatch 的 index 序号 # 在一些细节场景下需要有该信息区分一些资源的申请和管理。 self.microbatch_index: int = 0 - self.recurrent_state_context: Optional["RecurrentStateContext"] = None + self.ssm_run_config: Optional["SSMRunConfig"] = None self.b_logits_rows: Optional[torch.Tensor] = None self.logits_row_stride: int = 1 diff --git a/lightllm/common/recurrent_state.py b/lightllm/common/recurrent_state.py index 323c3e2b55..f133805edd 100644 --- a/lightllm/common/recurrent_state.py +++ b/lightllm/common/recurrent_state.py @@ -10,7 +10,7 @@ import torch from lightllm.common.state_cache_manager import LayerCache, LinearAttCacheConfig -from lightllm.common.recurrent_state_types import RecurrentStateContext, SSMUpdateBackend +from lightllm.common.recurrent_state_types import SSMRunConfig, SSMUpdateBackend if TYPE_CHECKING: from lightllm.common.basemodel.triton_kernel.linear_att.replayssm import ReplaySSMCache as _ReplaySSMCache @@ -119,12 +119,12 @@ def commit( b_mtp_index, accepted_index, verify_width, - batches: Optional[Sequence[Tuple[int, Optional[RecurrentStateContext]]]] = None, + batches: Optional[Sequence[Tuple[int, Optional[SSMRunConfig]]]] = None, ): """Commit accepted prefixes without requiring full-state materialization. ``batches`` partitions logical requests in the same order as the MTP - start locations. Each context belongs to the forward that wrote that + start locations. Each configuration belongs to the forward that wrote that partition's scratch; padding rows do not participate in acceptance. """ from lightllm.common.basemodel.triton_kernel.mtp_utils import linear_att_mtp_state_index_update @@ -150,12 +150,12 @@ def commit( return assert sum(size for size, _ in batches) == reqs.numel() start = 0 - for size, context in batches: + for size, run_config in batches: if size: self.updates.accept_updates( reqs[start : start + size], self.accepted_index, - run_config=context.config if context is not None else None, + run_config=run_config, ) start += size diff --git a/lightllm/common/recurrent_state_types.py b/lightllm/common/recurrent_state_types.py index 69e25227e1..276caed8d1 100644 --- a/lightllm/common/recurrent_state_types.py +++ b/lightllm/common/recurrent_state_types.py @@ -43,25 +43,6 @@ def __len__(self) -> int: return 3 -@dataclass(frozen=True, slots=True) -class RecurrentStateContext: - """Per-forward/graph commit metadata, passed unchanged by the runner. - - Both verify and compact reconstruction read the same immutable layout. - Publishing a context detaches it from the tuner's mutable dictionaries. - """ - - config: SSMRunConfig - - def __post_init__(self): - if not isinstance(self.config, SSMRunConfig): - raise TypeError("Use from_run_config() to freeze a mutable kernel configuration") - - @classmethod - def from_run_config(cls, config: Mapping[str, int]) -> "RecurrentStateContext": - return cls(SSMRunConfig.from_mapping(config)) - - class SSMUpdateBackend(Protocol): """Attention executes the updater directly; lifecycle stays with its owner. diff --git a/lightllm/server/router/model_infer/mode_backend/chunked_prefill/impl.py b/lightllm/server/router/model_infer/mode_backend/chunked_prefill/impl.py index 95f2c358b5..e1e4a507c2 100644 --- a/lightllm/server/router/model_infer/mode_backend/chunked_prefill/impl.py +++ b/lightllm/server/router/model_infer/mode_backend/chunked_prefill/impl.py @@ -115,7 +115,12 @@ def prefill_normal( event_pack.notify_post_handle_and_wait_pre_post_handle() with torch.cuda.stream(g_infer_context.get_overlap_stream()): - (_, next_token_ids_cpu, next_token_logprobs_cpu, next_token_ranks_cpu,) = self._sample_and_scatter_token( + ( + _, + next_token_ids_cpu, + next_token_logprobs_cpu, + next_token_ranks_cpu, + ) = self._sample_and_scatter_token( logits=model_output.logits, b_req_idx=model_input.b_req_idx, b_mtp_index=model_input.b_mtp_index, @@ -160,7 +165,12 @@ def decode_normal( event_pack.notify_post_handle_and_wait_pre_post_handle() with torch.cuda.stream(g_infer_context.get_overlap_stream()): - (_, next_token_ids_cpu, next_token_logprobs_cpu, next_token_ranks_cpu,) = self._sample_and_scatter_token( + ( + _, + next_token_ids_cpu, + next_token_logprobs_cpu, + next_token_ranks_cpu, + ) = self._sample_and_scatter_token( logits=model_output.logits, b_req_idx=model_input.b_req_idx, b_mtp_index=model_input.b_mtp_index, @@ -296,7 +306,7 @@ def decode_mtp( b_req_idx=model_input.b_req_idx, b_req_mtp_start_loc=b_req_mtp_start_loc, b_mtp_index=model_input.b_mtp_index, - recurrent_state_batches=((req_num, model_output.recurrent_state_context),), + recurrent_state_batches=((req_num, model_output.ssm_run_config),), ) accepted_index_cpu = g_pin_mem_manager.async_copy_from_gpu_tensor( key="accepted_index", diff --git a/lightllm/server/router/model_infer/mode_backend/dp_backend/impl.py b/lightllm/server/router/model_infer/mode_backend/dp_backend/impl.py index 95a7d12202..ca0c0d0231 100644 --- a/lightllm/server/router/model_infer/mode_backend/dp_backend/impl.py +++ b/lightllm/server/router/model_infer/mode_backend/dp_backend/impl.py @@ -525,7 +525,7 @@ def decode_mtp(self, event_pack: OverlapEventPack, decode_reqs: List[InferReq]): b_req_idx=model_input.b_req_idx, b_req_mtp_start_loc=b_req_mtp_start_loc, b_mtp_index=model_input.b_mtp_index, - recurrent_state_batches=((req_num, model_output.recurrent_state_context),), + recurrent_state_batches=((req_num, model_output.ssm_run_config),), ) accepted_index_cpu = g_pin_mem_manager.async_copy_from_gpu_tensor( key="accepted_index", @@ -801,8 +801,8 @@ def decode_overlap_mtp(self, event_pack: OverlapEventPack, decode_reqs: List[Inf b_req_mtp_start_loc=b_req_mtp_start_loc, b_mtp_index=b_mtp_index, recurrent_state_batches=( - (real_request_num0, model_output0.recurrent_state_context), - (real_request_num1, model_output1.recurrent_state_context), + (real_request_num0, model_output0.ssm_run_config), + (real_request_num1, model_output1.ssm_run_config), ), ) mtp_accept_len0 = mtp_accept_len[:real_request_num0] diff --git a/unit_tests/common/basemodel/test_draft_tail_graph.py b/unit_tests/common/basemodel/test_draft_tail_graph.py index e1aa933a4a..75bdf67955 100644 --- a/unit_tests/common/basemodel/test_draft_tail_graph.py +++ b/unit_tests/common/basemodel/test_draft_tail_graph.py @@ -63,7 +63,7 @@ def replay(): model._select_mem_indexes = lambda _: None model.decode_att_backend = SimpleNamespace( create_att_decode_state=lambda **kw: SimpleNamespace( - init_state=lambda: None, copy_for_decode_cuda_graph=lambda _: None + init_state=lambda: None, copy_for_decode_cuda_graph=lambda _: None, check_no_infer_state=lambda: None ) ) model.decode_att_backend1 = None diff --git a/unit_tests/common/basemodel/triton_kernel/linear_att/test_ssm_autotune.py b/unit_tests/common/basemodel/triton_kernel/linear_att/test_ssm_autotune.py index 419cb47e81..f88689f3df 100644 --- a/unit_tests/common/basemodel/triton_kernel/linear_att/test_ssm_autotune.py +++ b/unit_tests/common/basemodel/triton_kernel/linear_att/test_ssm_autotune.py @@ -7,7 +7,8 @@ from lightllm.common.basemodel.triton_kernel.linear_att import ssm_autotune as tuning from lightllm.common.basemodel.triton_kernel.linear_att.replayssm import ReplaySSMCache from lightllm.common.basemodel.triton_kernel.linear_att.replayssm_compact import CompactSSMCache -from lightllm.common.recurrent_state import RecurrentState, RecurrentStateContext +from lightllm.common.recurrent_state import RecurrentState +from lightllm.common.recurrent_state_types import SSMRunConfig from lightllm.common.kernel_config import KernelConfigs from lightllm.common import kernel_config as kernel_config_module from lightllm.common.triton_utils import autotuner as autotuner_module @@ -352,14 +353,14 @@ def forward(target): state = SimpleNamespace( hidden_collector=SimpleNamespace(finish_output=lambda **kw: ModelMtpOutputCollector()), prompt_logics=None, - recurrent_state_context=RecurrentStateContext.from_run_config(config), + ssm_run_config=SSMRunConfig.from_mapping(config), logits_row_stride=1, ) output = TpPartBaseModel._create_model_output( None, PostLayerOutput(logits=outputs[0].view(count * 4, -1)), state ) output = TpPartBaseModel._create_unpad_decode_model_output(None, output, len(ids) * 4) - assert dict(output.recurrent_state_context.config) == config + assert dict(output.ssm_run_config) == config buckets.append((graph, outputs, output, ref, args, reqs, cu, ids)) # Published contexts and captured launches must survive later changes to # the tuner's dictionaries, including when two graphs replay alternately. @@ -389,7 +390,7 @@ def forward(target): mtp, flags, 4, - recurrent_state_batches=tuple((len(b[-1]), b[2].recurrent_state_context) for b in buckets), + recurrent_state_batches=tuple((len(b[-1]), b[2].ssm_run_config) for b in buckets), ) for _, _, _, ref, _, reqs, _, ids in buckets: ref.accept_updates(reqs[:-1], manager.recurrent_state.accepted_index) diff --git a/unit_tests/common/test_recurrent_state_contract.py b/unit_tests/common/test_recurrent_state_contract.py index 82b701f0e7..b3d36a99ab 100644 --- a/unit_tests/common/test_recurrent_state_contract.py +++ b/unit_tests/common/test_recurrent_state_contract.py @@ -6,32 +6,30 @@ import torch from lightllm.common.recurrent_state import RecurrentState -from lightllm.common.recurrent_state_types import RecurrentStateContext +from lightllm.common.recurrent_state_types import SSMRunConfig from lightllm.common.state_cache_manager import Glm5NextCacheConfig, LinearAttCacheConfig from lightllm.server.core.objs.start_args_type import StartArgs -def test_published_context_is_detached_immutable_and_pickleable(): +def test_published_config_is_detached_immutable_and_pickleable(): tuning = {"BV": 8, "num_warps": 1} - context = RecurrentStateContext.from_run_config(tuning) + config = SSMRunConfig.from_mapping(tuning) tuning.update(BV=128, num_warps=4, num_stages=1) - assert dict(context.config) == {"BV": 8, "num_warps": 1, "num_stages": 3} + assert dict(config) == {"BV": 8, "num_warps": 1, "num_stages": 3} with pytest.raises(FrozenInstanceError): - context.config.BV = 128 + config.BV = 128 with pytest.raises(TypeError): - context.config["BV"] = 128 - with pytest.raises(TypeError, match="from_run_config"): - RecurrentStateContext(tuning) - restored = pickle.loads(pickle.dumps(context)) - assert restored == context - assert dict(restored.config) == dict(context.config) + config["BV"] = 128 + restored = pickle.loads(pickle.dumps(config)) + assert restored == config + assert dict(restored) == dict(config) -def test_context_rejects_unrepresented_kernel_options(): +def test_config_rejects_unrepresented_kernel_options(): with pytest.raises(ValueError, match="Unknown SSM kernel configuration fields"): - RecurrentStateContext.from_run_config({"BV": 32, "num_warps": 1, "future_layout_option": 2}) + SSMRunConfig.from_mapping({"BV": 32, "num_warps": 1, "future_layout_option": 2}) with pytest.raises(TypeError, match="must be an integer"): - RecurrentStateContext.from_run_config({"BV": [32], "num_warps": 1}) + SSMRunConfig.from_mapping({"BV": [32], "num_warps": 1}) @pytest.fixture diff --git a/unit_tests/models/test_post_layer_output.py b/unit_tests/models/test_post_layer_output.py index 43ff69e8c3..c8e081862d 100644 --- a/unit_tests/models/test_post_layer_output.py +++ b/unit_tests/models/test_post_layer_output.py @@ -291,7 +291,7 @@ def test_model_output_and_unpadding_preserve_candidate_mapping() -> None: logits = torch.randn(3, 2) ids = torch.tensor([[10, 20], [30, 40], [50, 60]]) state = SimpleNamespace( - hidden_collector=NoopHiddenCollector(), prompt_logics=None, recurrent_state_context=None, logits_row_stride=1 + hidden_collector=NoopHiddenCollector(), prompt_logics=None, ssm_run_config=None, logits_row_stride=1 ) output = model._create_model_output(PostLayerOutput(logits, ids), state) unpadded = model._create_unpad_decode_model_output(output, 2) From 12532af5b876ac21ce0763e5c52f6a9a90e59c44 Mon Sep 17 00:00:00 2001 From: sufubao Date: Mon, 19 Oct 2026 18:09:18 +0800 Subject: [PATCH 27/40] style: group PR imports and clarify SSM layer indices --- .../common/basemodel/attention/linear/gdn.py | 6 ++- .../common/basemodel/attention/linear/kda.py | 2 +- lightllm/common/basemodel/basemodel.py | 3 +- lightllm/common/basemodel/cuda_graph.py | 8 +-- lightllm/common/basemodel/infer_struct.py | 6 ++- .../triton_kernel/linear_att/replayssm.py | 52 ++++++++++--------- .../linear_att/replayssm_compact.py | 20 +++---- .../triton_kernel/linear_att/ssm_autotune.py | 5 +- .../post_process/apply_penalty_gpu_cache.py | 1 + .../kv_cache_mem_manager/mem_manager.py | 10 ++-- .../qwen3next_mem_manager.py | 4 +- lightllm/common/recurrent_state.py | 4 +- lightllm/common/recurrent_state_types.py | 2 +- .../common/state_cache_manager/linear_att.py | 4 +- .../llama/layer_infer/post_layer_infer.py | 4 +- .../layer_infer/transformer_layer_infer.py | 8 +-- lightllm/server/api_start.py | 1 + .../server/router/model_infer/infer_batch.py | 12 ++--- .../model_infer/mode_backend/base_backend.py | 8 +-- .../mode_backend/chunked_prefill/impl.py | 4 +- .../mode_backend/dp_backend/impl.py | 4 +- .../mode_backend/generic_post_process.py | 4 +- .../kernels/benchmark_replayssm_service.py | 1 + .../basemodel/attention/linear/test_gdn.py | 7 +-- .../linear_att/test_acceptance_fusion.py | 1 + unit_tests/models/glm5_next/test_mtp.py | 1 + unit_tests/server/test_ssm_state_mode.py | 1 + 27 files changed, 107 insertions(+), 76 deletions(-) diff --git a/lightllm/common/basemodel/attention/linear/gdn.py b/lightllm/common/basemodel/attention/linear/gdn.py index 283df27afe..2a756298f2 100644 --- a/lightllm/common/basemodel/attention/linear/gdn.py +++ b/lightllm/common/basemodel/attention/linear/gdn.py @@ -1,9 +1,11 @@ import dataclasses import weakref -import torch -from lightllm.common.recurrent_state_types import SSMRunConfig from abc import ABC, abstractmethod from typing import TYPE_CHECKING + +import torch + +from lightllm.common.recurrent_state_types import SSMRunConfig from ..base_att import BaseAttBackend, BasePrefillAttState, BaseDecodeAttState, AttControl from lightllm.common.req_manager import ReqManagerForMamba from lightllm.utils.envs_utils import get_env_start_args diff --git a/lightllm/common/basemodel/attention/linear/kda.py b/lightllm/common/basemodel/attention/linear/kda.py index a2fe2288c0..6285a2b384 100644 --- a/lightllm/common/basemodel/attention/linear/kda.py +++ b/lightllm/common/basemodel/attention/linear/kda.py @@ -9,8 +9,8 @@ from typing import TYPE_CHECKING import torch -from lightllm.common.recurrent_state_types import SSMRunConfig +from lightllm.common.recurrent_state_types import SSMRunConfig from lightllm.common.basemodel.attention.base_att import ( AttControl, BaseAttBackend, diff --git a/lightllm/common/basemodel/basemodel.py b/lightllm/common/basemodel/basemodel.py index 0d43e6d9dd..eca299f113 100755 --- a/lightllm/common/basemodel/basemodel.py +++ b/lightllm/common/basemodel/basemodel.py @@ -5,10 +5,11 @@ import copy import math import json +from typing import final, List + import torch import torch.nn.functional as F import triton -from typing import final, List from tqdm import tqdm from lightllm.common.basemodel.layer_weights.hf_load_utils import load_hf_weights diff --git a/lightllm/common/basemodel/cuda_graph.py b/lightllm/common/basemodel/cuda_graph.py index fa236757be..5981a5bd18 100644 --- a/lightllm/common/basemodel/cuda_graph.py +++ b/lightllm/common/basemodel/cuda_graph.py @@ -1,11 +1,13 @@ import os -import torch -import torch.distributed as dist import copy import bisect import math -import triton from typing import Optional + +import torch +import torch.distributed as dist +import triton + from lightllm.utils.log_utils import init_logger from lightllm.utils.envs_utils import get_env_start_args from lightllm.distributed import dist_group_manager diff --git a/lightllm/common/basemodel/infer_struct.py b/lightllm/common/basemodel/infer_struct.py index 0f94a34a10..acfd371c95 100755 --- a/lightllm/common/basemodel/infer_struct.py +++ b/lightllm/common/basemodel/infer_struct.py @@ -1,10 +1,12 @@ +import collections +from typing import TYPE_CHECKING, Tuple, Any, Optional, List + import torch import triton -import collections + from lightllm.common.kv_cache_mem_manager import MemoryManager from lightllm.common.req_manager import ReqManager from lightllm.distributed import CustomProcessGroup -from typing import TYPE_CHECKING, Tuple, Any, Optional, List from .triton_kernel.gen_prefill_params import gen_prefill_params from .triton_kernel.gen_decode_params import gen_decode_params from .triton_kernel.multimodal_emb import mark_multimodal_obj diff --git a/lightllm/common/basemodel/triton_kernel/linear_att/replayssm.py b/lightllm/common/basemodel/triton_kernel/linear_att/replayssm.py index c1f4ca271c..75bf462808 100644 --- a/lightllm/common/basemodel/triton_kernel/linear_att/replayssm.py +++ b/lightllm/common/basemodel/triton_kernel/linear_att/replayssm.py @@ -131,7 +131,7 @@ def _fold_history( slot_work = work // tiles seq = tl.load(Active + slot_work % count) layer_hv = slot_work // count - layer, hv = (layer_hv // HV).to(tl.int64), layer_hv % HV + linear_layer_idx, hv = (layer_hv // HV).to(tl.int64), layer_hv % HV req = tl.load(Reqs + seq).to(tl.int64) cursor = tl.load(Cursors + seq) n = cursor % (2 * L) @@ -140,21 +140,21 @@ def _fold_history( vv = iv * BV + tl.arange(0, BV) slot = req * HV + hv key_slot = req * KH + hv // (HV // KH) - sp = State + layer * SS + slot * K * V + kk[:, None] * V + vv[None, :] + sp = State + linear_layer_idx * SS + slot * K * V + kk[:, None] * V + vv[None, :] state = tl.load(sp, (kk[:, None] < K) & (vv[None, :] < V), 0).to(tl.float32) if DERIVED: tl.static_assert(not KDA) ll = tl.arange(0, max(16, L)) - g = tl.load(Gates + layer * GS + slot * (2 * L) + base + ll, ll < n, 0).to(tl.float32) + g = tl.load(Gates + linear_layer_idx * GS + slot * (2 * L) + base + ll, ll < n, 0).to(tl.float32) prefix = tl.cumsum(g, 0) total_g = tl.sum(g, 0) keys = tl.load( - Keys + layer * RS + (key_slot * (2 * L) + base + ll[None, :]) * K + kk[:, None], + Keys + linear_layer_idx * RS + (key_slot * (2 * L) + base + ll[None, :]) * K + kk[:, None], (kk[:, None] < K) & (ll[None, :] < n), 0, ).to(tl.float32) delta = tl.load( - Deltas + layer * VS + (slot * (2 * L) + base + ll[:, None]) * V + vv[None, :], + Deltas + linear_layer_idx * VS + (slot * (2 * L) + base + ll[:, None]) * V + vv[None, :], (ll[:, None] < n) & (vv[None, :] < V), 0, ).to(tl.float32) @@ -162,15 +162,19 @@ def _fold_history( state = state * tl.exp(total_g) + tl.dot(keys, delta, input_precision="tf32x3") else: for j in range(n): - raw_k = tl.load(RawKeys + layer * RS + (key_slot * (2 * L) + base + j) * K + kk, kk < K, 0).to( + raw_k = tl.load( + RawKeys + linear_layer_idx * RS + (key_slot * (2 * L) + base + j) * K + kk, kk < K, 0 + ).to(tl.float32) + raw_v = tl.load(RawValues + linear_layer_idx * VS + (slot * (2 * L) + base + j) * V + vv, vv < V, 0).to( tl.float32 ) - raw_v = tl.load(RawValues + layer * VS + (slot * (2 * L) + base + j) * V + vv, vv < V, 0).to(tl.float32) if KDA: - g = tl.load(Gates + layer * GS + (slot * (2 * L) + base + j) * K + kk, kk < K, 0).to(tl.float32) + g = tl.load(Gates + linear_layer_idx * GS + (slot * (2 * L) + base + j) * K + kk, kk < K, 0).to( + tl.float32 + ) else: - g = tl.load(Gates + layer * GS + slot * (2 * L) + base + j).to(tl.float32) - beta = tl.load(Betas + layer * BS + slot * (2 * L) + base + j).to(tl.float32) + g = tl.load(Gates + linear_layer_idx * GS + slot * (2 * L) + base + j).to(tl.float32) + beta = tl.load(Betas + linear_layer_idx * BS + slot * (2 * L) + base + j).to(tl.float32) state = _replay_update(state, raw_k, raw_v, g, beta, KDA) # Start the new history from the same rounded checkpoint future calls load. state = state.to(State.dtype.element_ty).to(tl.float32) @@ -284,7 +288,7 @@ def _replay( qk /= tl.sqrt(tl.sum(qk * qk, axis=1) + 1.0e-6)[:, None] else: qk *= tl.rsqrt(tl.sum(qk * qk, axis=1) + 1.0e-6)[:, None] - qk *= tl.where(tt < BW, K ** -0.5, 1.0)[:, None] + qk *= tl.where(tt < BW, K**-0.5, 1.0)[:, None] if KDA: history = ll < n gates = tl.load( @@ -467,7 +471,7 @@ def _replay( 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) + 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) if KDA: x = tl.load(A + t * SA + hv * K + kk, kk < K, 0).to(tl.float32) + bias @@ -554,9 +558,9 @@ def _materialize( base = (cursor // (2 * L)) * L if n == 0 and not SNAPSHOT: return - layer, head = lh // HV, lh % HV - slot = (layer * SLOTS + req) * HV + head - key_slot = (layer * SLOTS + req) * KH + head // (HV // KH) + linear_layer_idx, head = lh // HV, lh % HV + slot = (linear_layer_idx * SLOTS + req) * HV + head + key_slot = (linear_layer_idx * SLOTS + req) * KH + head // (HV // KH) kk = tl.arange(0, BK) vv = iv * BV + tl.arange(0, BV) sp = State + slot * K * V + kk[:, None] * V + vv[None, :] @@ -573,7 +577,7 @@ def _materialize( state *= tl.exp(g)[:, None] if KDA else tl.exp(g) d = beta * (raw_v - tl.sum(state * raw_k[:, None], 0)) state += raw_k[:, None] * d[None, :] - out_slot = layer * HV + head if SNAPSHOT else slot + out_slot = linear_layer_idx * HV + head if SNAPSHOT else slot out = Output + out_slot * K * V + kk[:, None] * V + vv[None, :] tl.store(out, state, (kk[:, None] < K) & (vv[None, :] < V)) @@ -772,7 +776,7 @@ def get_run_config(self, q, k, v, a, b, a_log, bias, cu_seqlens=None): return self.run_config or {"BV": 128, "num_warps": 2, "num_stages": 1} return configure_cache(self, "replay", q, k, v, a, b, a_log, bias, cu_seqlens) - def forward(self, layer, q, k, v, a, b, a_log, bias, reqs, cu_seqlens=None, run_config=None): + def forward(self, linear_layer_idx, q, k, v, a, b, a_log, bias, reqs, cu_seqlens=None, run_config=None): config = run_config or self.get_run_config(q, k, v, a, b, a_log, bias, cu_seqlens) hv, kd, vd = self.state.shape[-3:] assert self.num_key_heads in (q.shape[-2], hv) @@ -790,13 +794,13 @@ def forward(self, layer, q, k, v, a, b, a_log, bias, reqs, cu_seqlens=None, run_ 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], + self.state[linear_layer_idx], + self.keys[linear_layer_idx], + self.deltas[linear_layer_idx], + self.gates[linear_layer_idx], + self.raw_keys[linear_layer_idx], + self.raw_values[linear_layer_idx], + self.betas[linear_layer_idx], reqs, self.forward_cursors, cu_seqlens, diff --git a/lightllm/common/basemodel/triton_kernel/linear_att/replayssm_compact.py b/lightllm/common/basemodel/triton_kernel/linear_att/replayssm_compact.py index 8f55269fd9..9ca2b75ef6 100644 --- a/lightllm/common/basemodel/triton_kernel/linear_att/replayssm_compact.py +++ b/lightllm/common/basemodel/triton_kernel/linear_att/replayssm_compact.py @@ -51,7 +51,7 @@ def _compact( ): 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 + linear_layer_idx, hv = lh // HV, lh % HV if COMMIT: start = 0 end = tl.load(Accepted + req) + 1 if req != HOLD else 0 @@ -64,8 +64,8 @@ def _compact( for t in range(start, end): tl.store(Out + (t * HV + hv) * V + vv, 0, vv < V) return - slot = (layer * SLOTS + req) * HV + hv - key_slot = (layer * SLOTS + req) * KH + hv // (HV // KH) + slot = (linear_layer_idx * SLOTS + req) * HV + hv + key_slot = (linear_layer_idx * SLOTS + req) * KH + hv // (HV // KH) 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): @@ -108,7 +108,7 @@ def _compact( state += k[:, None] * d[None, :] if not COMMIT: q = tl.load(Q + t * SQ + h * K + kk, kk < K, 0).to(tl.float32) - q = q / tl.sqrt(tl.sum(q * q) + 1.0e-6) * (K ** -0.5) + 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: @@ -159,7 +159,7 @@ def snapshot_accepted_state(self, req_idx): def get_run_config(self, q, k, v, a, b, a_log, bias, cu_seqlens=None): return configure_cache(self, "gdn", q, k, v, a, b, a_log, bias, cu_seqlens) - def forward(self, layer, q, k, v, a, b, a_log, bias, reqs, cu_seqlens=None, run_config=None): + def forward(self, linear_layer_idx, q, k, v, a, b, a_log, bias, reqs, cu_seqlens=None, run_config=None): assert cu_seqlens is not None, "compact replay is only used for speculative verify" config = run_config or self.get_run_config(q, k, v, a, b, a_log, bias, cu_seqlens) _, slots, hv, kd, vd = self.state.shape @@ -175,11 +175,11 @@ def forward(self, layer, q, k, v, a, b, a_log, bias, reqs, cu_seqlens=None, run_ b, a_log, bias, - self.state[layer], - self.keys[layer], - self.values[layer], - self.decays[layer], - self.betas[layer], + self.state[linear_layer_idx], + self.keys[linear_layer_idx], + self.values[linear_layer_idx], + self.decays[linear_layer_idx], + self.betas[linear_layer_idx], reqs, cu_seqlens, None, diff --git a/lightllm/common/basemodel/triton_kernel/linear_att/ssm_autotune.py b/lightllm/common/basemodel/triton_kernel/linear_att/ssm_autotune.py index cf623a7a1e..6546764134 100644 --- a/lightllm/common/basemodel/triton_kernel/linear_att/ssm_autotune.py +++ b/lightllm/common/basemodel/triton_kernel/linear_att/ssm_autotune.py @@ -1,6 +1,7 @@ """Tune a complete SSM update cycle without touching the serving state pool.""" import ast + import torch import triton @@ -139,8 +140,8 @@ def select_config(cache, mode, q, k, v, a, b, a_log, bias, cu_seqlens=None, work rounds = 2 # Partial acceptance and full acceptance, both across all layers. for step in range(rounds): cache.prepare_decode(reqs, cu_seqlens) - for layer in range(cache.state.shape[0]): - cache.forward(layer, q, k, v, a, b, a_log, bias, reqs, cu_seqlens) + for linear_layer_idx in range(cache.state.shape[0]): + cache.forward(linear_layer_idx, q, k, v, a, b, a_log, bias, reqs, cu_seqlens) cache.accept_updates(reqs, partial if mode != "replay" and step == 0 else accepted) if mode == "replay" and step == rounds - 2: cache._materialize_accepted_state(reqs[:1], snapshot, snapshot=True) diff --git a/lightllm/common/basemodel/triton_kernel/post_process/apply_penalty_gpu_cache.py b/lightllm/common/basemodel/triton_kernel/post_process/apply_penalty_gpu_cache.py index 9d70ea034e..f0fd2a4c78 100644 --- a/lightllm/common/basemodel/triton_kernel/post_process/apply_penalty_gpu_cache.py +++ b/lightllm/common/basemodel/triton_kernel/post_process/apply_penalty_gpu_cache.py @@ -3,6 +3,7 @@ import triton.language as tl import torch.nn.functional as F import numpy as np + from lightllm.common.req_manager import ReqSamplingParamsManager diff --git a/lightllm/common/kv_cache_mem_manager/mem_manager.py b/lightllm/common/kv_cache_mem_manager/mem_manager.py index eb8aee4e96..2d6e6e0105 100755 --- a/lightllm/common/kv_cache_mem_manager/mem_manager.py +++ b/lightllm/common/kv_cache_mem_manager/mem_manager.py @@ -1,10 +1,14 @@ import math import os import re +from typing import List, Tuple, Any, Union +from multiprocessing.reduction import ForkingPickler + import torch import torch.distributed as dist import torch.multiprocessing as mp -from typing import List, Tuple, Any, Union +from filelock import FileLock + from lightllm.utils.log_utils import init_logger from lightllm.server.router.dynamic_prompt.shared_arr import SharedInt from .allocator import KvCacheAllocator @@ -20,8 +24,6 @@ from lightllm.common.kv_trans_kernel.nixl_kv_trans import page_io from lightllm.utils.device_utils import kv_trans_use_p2p from lightllm.utils.shm_utils import create_or_link_shm -from multiprocessing.reduction import ForkingPickler -from filelock import FileLock from .operator import BaseMemManagerOperator, NormalMemOperator logger = init_logger(__name__) @@ -97,7 +99,7 @@ def profile_size(self, mem_fraction): available_memory = get_available_gpu_memory(world_size) - get_total_gpu_memory() * (1 - mem_fraction) cell_size = self.get_cell_size() pd_kv_move_buffer_size = self.get_pd_kv_move_buffer_size() - available_memory_bytes = available_memory * 1024 ** 3 - pd_kv_move_buffer_size + available_memory_bytes = available_memory * 1024**3 - pd_kv_move_buffer_size self.size = int(available_memory_bytes / cell_size) if world_size > 1: tensor = torch.tensor(self.size, dtype=torch.int64, device=f"cuda:{get_current_device_id()}") 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 e5365ea449..315f9dfcda 100644 --- a/lightllm/common/kv_cache_mem_manager/qwen3next_mem_manager.py +++ b/lightllm/common/kv_cache_mem_manager/qwen3next_mem_manager.py @@ -1,5 +1,8 @@ +from typing import Tuple, Any, List + import torch import triton + from lightllm.utils.log_utils import init_logger from lightllm.common.kv_cache_mem_manager.mem_manager import MemoryManager from lightllm.utils.envs_utils import get_env_start_args @@ -11,7 +14,6 @@ FP8StaticPerTensorQuantMemOperator, LinearAttMemOperator, ) -from typing import Tuple, Any, List logger = init_logger(__name__) diff --git a/lightllm/common/recurrent_state.py b/lightllm/common/recurrent_state.py index f133805edd..d357c1ed06 100644 --- a/lightllm/common/recurrent_state.py +++ b/lightllm/common/recurrent_state.py @@ -230,5 +230,5 @@ def export(self, req_idx, conv, ssm): ssm.copy_(snapshot.ssm, non_blocking=True) def layer_buffers(self, layer_idx_in_all): - layer = self.linear_config.get_linear_layer_index(layer_idx_in_all) - return self.conv_state.buffer[layer], self.ssm_state.buffer[layer] + linear_layer_idx = self.linear_config.get_linear_layer_index(layer_idx_in_all) + return self.conv_state.buffer[linear_layer_idx], self.ssm_state.buffer[linear_layer_idx] diff --git a/lightllm/common/recurrent_state_types.py b/lightllm/common/recurrent_state_types.py index 276caed8d1..ceb13e74f7 100644 --- a/lightllm/common/recurrent_state_types.py +++ b/lightllm/common/recurrent_state_types.py @@ -66,7 +66,7 @@ def get_run_config( def forward( self, - layer: int, + linear_layer_idx: int, q: torch.Tensor, k: torch.Tensor, v: torch.Tensor, diff --git a/lightllm/common/state_cache_manager/linear_att.py b/lightllm/common/state_cache_manager/linear_att.py index e533537deb..12d262e608 100644 --- a/lightllm/common/state_cache_manager/linear_att.py +++ b/lightllm/common/state_cache_manager/linear_att.py @@ -1,11 +1,11 @@ -import torch import dataclasses + +import torch import triton from lightllm.utils.envs_utils import get_added_mtp_kv_layer_num, get_env_start_args from lightllm.utils.log_utils import init_logger from lightllm.utils.torch_dtype_utils import get_torch_dtype - from .base import StateCacheManager from .layer_cache import LayerCache diff --git a/lightllm/models/llama/layer_infer/post_layer_infer.py b/lightllm/models/llama/layer_infer/post_layer_infer.py index dbdb8ee0ba..c186877eea 100644 --- a/lightllm/models/llama/layer_infer/post_layer_infer.py +++ b/lightllm/models/llama/layer_infer/post_layer_infer.py @@ -1,9 +1,11 @@ import os + import torch -from lightllm.common.basemodel.batch_objs import PostLayerOutput import torch.functional as F import torch.distributed as dist import numpy as np + +from lightllm.common.basemodel.batch_objs import PostLayerOutput from lightllm.common.basemodel.layer_weights.base_layer_weight import BaseLayerWeight from lightllm.models.llama.layer_weights.pre_and_post_layer_weight import LlamaPreAndPostLayerWeight from lightllm.models.llama.infer_struct import LlamaInferStateInfo diff --git a/lightllm/models/qwen3next/layer_infer/transformer_layer_infer.py b/lightllm/models/qwen3next/layer_infer/transformer_layer_infer.py index 4ef9d0b53c..38d558782b 100644 --- a/lightllm/models/qwen3next/layer_infer/transformer_layer_infer.py +++ b/lightllm/models/qwen3next/layer_infer/transformer_layer_infer.py @@ -1,6 +1,9 @@ -import torch +from typing import Tuple +from functools import partial +import torch import torch.distributed as dist + from lightllm.models.qwen3next.layer_weights.transformer_layer_weight import ( Qwen3NextTransformerLayerWeight, ) @@ -10,12 +13,10 @@ from lightllm.utils.tensor_utils import tensor_to_no_ref_tensor from lightllm.common.kv_cache_mem_manager import Qwen3NextMemManager from lightllm.common.basemodel.attention.base_att import AttControl -from typing import Tuple from lightllm.models.qwen3next.triton_kernel.shared_expert_gate import sigmoid_mul_ from lightllm.distributed import all_reduce from lightllm.models.llama.triton_kernel.rotary_emb import rotary_emb_fwd from lightllm.utils.envs_utils import get_env_start_args -from functools import partial logger = init_logger(__name__) @@ -184,7 +185,6 @@ def _linear_prefill_cuda_graph_wrapper( infer_state: Qwen3NextInferStateInfo, layer_weight: Qwen3NextTransformerLayerWeight, ) -> Tuple[torch.Tensor, torch.Tensor]: - backend = infer_state.prefill_att_state1.backend mixed_qkvzba = mixed_qkvzba.contiguous() diff --git a/lightllm/server/api_start.py b/lightllm/server/api_start.py index e7e118fc25..ef26a09ed0 100644 --- a/lightllm/server/api_start.py +++ b/lightllm/server/api_start.py @@ -3,6 +3,7 @@ import uuid import subprocess import math + from lightllm.utils.start_utils import process_manager from .metrics.manager import start_metric_manager from .embed_cache.manager import start_cache_manager diff --git a/lightllm/server/router/model_infer/infer_batch.py b/lightllm/server/router/model_infer/infer_batch.py index e423a0e7c7..8527a2261d 100644 --- a/lightllm/server/router/model_infer/infer_batch.py +++ b/lightllm/server/router/model_infer/infer_batch.py @@ -1,13 +1,14 @@ import enum -import torch -import torch.distributed as dist -import numpy as np import collections import pickle - -from sortedcontainers import SortedDict from dataclasses import dataclass, field from typing import TYPE_CHECKING, List, Dict, Tuple, Optional, Callable, Any, Union + +import torch +import torch.distributed as dist +import numpy as np +from sortedcontainers import SortedDict + from lightllm.common.req_manager import ReqManager, HybridAttentionReqManager from lightllm.utils.infer_utils import mark_start, mark_end from lightllm.server.core.objs import Req, SamplingParams, FinishStatus, ShmReqManager @@ -371,7 +372,6 @@ def filter_reqs(self, finished_reqs: List["InferReq"]): @torch.no_grad() def pause_reqs(self, pause_reqs: List["InferReq"], is_master_in_dp: bool): if pause_reqs: - free_token_index = [] for req in pause_reqs: if self.args.diverse_mode: diff --git a/lightllm/server/router/model_infer/mode_backend/base_backend.py b/lightllm/server/router/model_infer/mode_backend/base_backend.py index 1ec27484e4..e27ef23b7c 100644 --- a/lightllm/server/router/model_infer/mode_backend/base_backend.py +++ b/lightllm/server/router/model_infer/mode_backend/base_backend.py @@ -1,12 +1,13 @@ import os +import time +import threading +from typing import List, Tuple, Callable, Optional, Union import numpy as np import torch -import time -import threading import torch.distributed as dist -from typing import List, Tuple, Callable, Optional, Union from transformers.configuration_utils import PretrainedConfig + from lightllm.utils.infer_utils import set_random_seed from lightllm.utils.log_utils import init_logger from lightllm.models import get_draft_model_class, get_model @@ -722,7 +723,6 @@ def _get_classed_reqs( can_alloc_token_num = g_infer_context.get_can_alloc_token_num() for req_obj in ready_reqs: - if req_obj.filter_mark: finished_reqs.append(req_obj) continue diff --git a/lightllm/server/router/model_infer/mode_backend/chunked_prefill/impl.py b/lightllm/server/router/model_infer/mode_backend/chunked_prefill/impl.py index e1e4a507c2..8ad3393f5f 100644 --- a/lightllm/server/router/model_infer/mode_backend/chunked_prefill/impl.py +++ b/lightllm/server/router/model_infer/mode_backend/chunked_prefill/impl.py @@ -1,6 +1,8 @@ -import torch import time from typing import List + +import torch + from lightllm.common.basemodel.triton_kernel.mtp_utils import gen_b_req_mtp_start_loc from lightllm.server.router.model_infer.mode_backend.base_backend import ModeBackend from lightllm.server.router.model_infer.mode_backend.overlap_events import OverlapEventPack diff --git a/lightllm/server/router/model_infer/mode_backend/dp_backend/impl.py b/lightllm/server/router/model_infer/mode_backend/dp_backend/impl.py index ca0c0d0231..8f4cfe42bf 100644 --- a/lightllm/server/router/model_infer/mode_backend/dp_backend/impl.py +++ b/lightllm/server/router/model_infer/mode_backend/dp_backend/impl.py @@ -1,6 +1,8 @@ -import torch import time from typing import List, Tuple + +import torch + from lightllm.common.basemodel.triton_kernel.mtp_utils import gen_b_req_mtp_start_loc from lightllm.server.router.model_infer.mode_backend.base_backend import ModeBackend from lightllm.common.basemodel.batch_objs import ModelOutput, ModelInput diff --git a/lightllm/server/router/model_infer/mode_backend/generic_post_process.py b/lightllm/server/router/model_infer/mode_backend/generic_post_process.py index 71023bd801..68c6822645 100644 --- a/lightllm/server/router/model_infer/mode_backend/generic_post_process.py +++ b/lightllm/server/router/model_infer/mode_backend/generic_post_process.py @@ -1,5 +1,7 @@ -import torch from typing import List, Optional, Tuple + +import torch + from lightllm.common.basemodel.triton_kernel.post_process.apply_penalty import apply_penalty from lightllm.common.basemodel.triton_kernel.post_process.apply_penalty_gpu_cache import apply_penalty_gpu_cache from lightllm.common.basemodel.triton_kernel.post_process.apply_invalid_token import apply_invalid_token_ids diff --git a/test/benchmark/kernels/benchmark_replayssm_service.py b/test/benchmark/kernels/benchmark_replayssm_service.py index ea65431ad4..08a694a684 100644 --- a/test/benchmark/kernels/benchmark_replayssm_service.py +++ b/test/benchmark/kernels/benchmark_replayssm_service.py @@ -10,6 +10,7 @@ import time import threading from pathlib import Path + import requests from transformers import AutoTokenizer diff --git a/unit_tests/common/basemodel/attention/linear/test_gdn.py b/unit_tests/common/basemodel/attention/linear/test_gdn.py index ab6f9d0483..c5bcd708e8 100644 --- a/unit_tests/common/basemodel/attention/linear/test_gdn.py +++ b/unit_tests/common/basemodel/attention/linear/test_gdn.py @@ -165,7 +165,7 @@ def build_params(**kwargs): def test_decode_state_passes_actual_verify_lengths_to_ssm_cache(monkeypatch): - state = _create_decode_state( + state, infer_state = _create_decode_state( draft_step=3, dynamic_layout=True, b_req_idx=torch.tensor([0, 0, 1, 1, 1]), @@ -210,8 +210,9 @@ def test_linear_backend_checks_req_manager_at_construction(monkeypatch, manager_ else: backend.__init__(model) assert backend.req_manager is manager - assert backend.create_att_prefill_state(None).backend.req_manager is manager - assert backend.create_att_decode_state(None).backend.req_manager is manager + infer_state = InferStateInfo() + assert backend.create_att_prefill_state(infer_state).backend.req_manager is manager + assert backend.create_att_decode_state(infer_state).backend.req_manager is manager @pytest.fixture diff --git a/unit_tests/common/basemodel/triton_kernel/linear_att/test_acceptance_fusion.py b/unit_tests/common/basemodel/triton_kernel/linear_att/test_acceptance_fusion.py index 83bfa18c5a..f431878851 100644 --- a/unit_tests/common/basemodel/triton_kernel/linear_att/test_acceptance_fusion.py +++ b/unit_tests/common/basemodel/triton_kernel/linear_att/test_acceptance_fusion.py @@ -1,6 +1,7 @@ import pytest import torch import triton + from lightllm.common.basemodel.triton_kernel.mtp_utils import linear_att_mtp_state_index_update as update from lightllm.common.basemodel.triton_kernel.linear_att.replayssm import ReplaySSMCache from lightllm.common.recurrent_state import RecurrentState diff --git a/unit_tests/models/glm5_next/test_mtp.py b/unit_tests/models/glm5_next/test_mtp.py index 75c8df2260..e88e12cfee 100644 --- a/unit_tests/models/glm5_next/test_mtp.py +++ b/unit_tests/models/glm5_next/test_mtp.py @@ -1,4 +1,5 @@ import weakref + import pytest import torch diff --git a/unit_tests/server/test_ssm_state_mode.py b/unit_tests/server/test_ssm_state_mode.py index 7f1aad5e74..58dee5938c 100644 --- a/unit_tests/server/test_ssm_state_mode.py +++ b/unit_tests/server/test_ssm_state_mode.py @@ -1,4 +1,5 @@ import json + import pytest from lightllm.server.api_cli import make_argument_parser From 0f663cd62ddeb7067d8319cc0b6b4cf1e7c09c23 Mon Sep 17 00:00:00 2001 From: sufubao Date: Tue, 20 Oct 2026 08:59:51 +0800 Subject: [PATCH 28/40] fix: pass ReplaySSM tile dimensions as compile-time constants --- .../triton_kernel/linear_att/replayssm.py | 23 +++++++++++++------ 1 file changed, 16 insertions(+), 7 deletions(-) diff --git a/lightllm/common/basemodel/triton_kernel/linear_att/replayssm.py b/lightllm/common/basemodel/triton_kernel/linear_att/replayssm.py index 75bf462808..f47d69e876 100644 --- a/lightllm/common/basemodel/triton_kernel/linear_att/replayssm.py +++ b/lightllm/common/basemodel/triton_kernel/linear_att/replayssm.py @@ -121,6 +121,7 @@ def _fold_history( Keys, Deltas, DERIVED: tl.constexpr, + BL: tl.constexpr, ): count = tl.load(Active + N) work = tl.program_id(0) @@ -144,7 +145,7 @@ def _fold_history( state = tl.load(sp, (kk[:, None] < K) & (vv[None, :] < V), 0).to(tl.float32) if DERIVED: tl.static_assert(not KDA) - ll = tl.arange(0, max(16, L)) + ll = tl.arange(0, BL) g = tl.load(Gates + linear_layer_idx * GS + slot * (2 * L) + base + ll, ll < n, 0).to(tl.float32) prefix = tl.cumsum(g, 0) total_g = tl.sum(g, 0) @@ -223,6 +224,8 @@ def _replay( KDA: tl.constexpr, LOWER: tl.constexpr, PARALLEL_SOLVE: tl.constexpr, + BQK: tl.constexpr, + BR: tl.constexpr, ): """Output reconstruction with an exact raw-input recurrence at checkpoint folds.""" iv, seq, hv = tl.program_id(0), tl.program_id(1), tl.program_id(2) @@ -269,7 +272,7 @@ def _replay( tl.bfloat16 if State.dtype.element_ty == tl.bfloat16 and Q.dtype.element_ty == tl.bfloat16 else tl.float32 ) projection_precision: tl.constexpr = "tf32" if Q.dtype.element_ty == tl.bfloat16 else "tf32x3" - tt = tl.arange(0, max(16, 2 * BW)) + tt = tl.arange(0, BQK) tokens = start + tt % BW valid = (tt < 2 * BW) & (tokens < end) h = hv // (HV // H) @@ -326,7 +329,7 @@ def _replay( if PRECOMPUTE_STATE and not KDA: # Project the whole verify window against both accepted and new keys. # Only the small corrected-value solve below depends on token order. - jj = tl.arange(0, triton.next_power_of_2(max(16, L + BW))) + jj = tl.arange(0, BR) fresh_tokens = start + jj - L old = (jj < L) & (jj < n) fresh = (jj >= L) & (jj < L + end - start) @@ -370,16 +373,16 @@ def _replay( projected_history = projections + tl.dot( coefficients.to(projection_dtype), deltas.to(projection_dtype), input_precision=projection_precision ) - wi = tl.arange(0, max(16, 2 * BW)) + wi = tl.arange(0, BQK) active = (wi < BW) & (start + wi < end) key_rows = tl.gather( projected_history, tl.broadcast_to((BW + wi % BW)[:, None], projected_history.shape), 0 ) - fresh_values = tl.gather(values, tl.broadcast_to((L + wi % BW)[:, None], (max(16, 2 * BW), BV)), 0) + fresh_values = tl.gather(values, tl.broadcast_to((L + wi % BW)[:, None], (BQK, BV)), 0) beta_rows = tl.sigmoid(tl.load(B + (start + wi) * SB + hv, active, 0).to(tl.float32)) rhs = tl.where(active[:, None], beta_rows[:, None] * (fresh_values - key_rows), 0.0) cross = tl.gather(coefficients, tl.broadcast_to((BW + wi % BW)[:, None], coefficients.shape), 0) - cross = tl.gather(cross, tl.broadcast_to((L + wi % BW)[None, :], (max(16, 2 * BW), max(16, 2 * BW))), 1) + cross = tl.gather(cross, tl.broadcast_to((L + wi % BW)[None, :], (BQK, BQK)), 1) lower = active[:, None] & active[None, :] & (wi[:, None] > wi[None, :]) power = tl.where(lower, -beta_rows[:, None] * cross, 0.0) inverse = (wi[:, None] == wi[None, :]).to(tl.float32) @@ -711,6 +714,7 @@ def prepare_decode(self, reqs, cu_seqlens=None): self.keys, self.deltas, self.fold_method == "derived", + BL=max(16, self.capacity), num_warps=2, num_stages=1, ) @@ -786,6 +790,9 @@ def forward(self, linear_layer_idx, q, k, v, a, b, a_log, bias, reqs, cu_seqlens precompute_state = self.verify_width > 1 and bv >= 16 and self.projection_mode in ("precompute", "parallel") if self.projection_mode == "parallel": assert precompute_state, "Parallel ReplaySSM requires BV >= 16" + # Resolve tile dimensions in Python: bare max() in Triton JIT code + # lowers to a tensor operation, which cannot define static shapes. + bw = triton.next_power_of_2(self.verify_width) _replay[(triton.cdiv(vd, bv), reqs.numel(), hv)]( q, k, @@ -821,11 +828,13 @@ def forward(self, linear_layer_idx, q, k, v, a, b, a_log, bias, reqs, cu_seqlens bv, cu_seqlens is not None, precompute_state, - triton.next_power_of_2(self.verify_width), + bw, reqs.numel() > 12 and self.projection_mode != "parallel", self.kda, self.lower_bound, self.projection_mode == "parallel", + BQK=max(16, 2 * bw), + BR=triton.next_power_of_2(max(16, self.capacity + bw)), num_warps=config["num_warps"], num_stages=config.get("num_stages", 3), ) From db380b5623217246cda7e71e0e14eb20e7dbd818 Mon Sep 17 00:00:00 2001 From: sufubao Date: Tue, 20 Oct 2026 11:54:34 +0800 Subject: [PATCH 29/40] style: match pinned Black formatting --- .../triton_kernel/linear_att/replayssm.py | 4 ++-- .../triton_kernel/linear_att/replayssm_compact.py | 2 +- .../common/kv_cache_mem_manager/mem_manager.py | 2 +- .../mode_backend/chunked_prefill/impl.py | 14 ++------------ 4 files changed, 6 insertions(+), 16 deletions(-) diff --git a/lightllm/common/basemodel/triton_kernel/linear_att/replayssm.py b/lightllm/common/basemodel/triton_kernel/linear_att/replayssm.py index f47d69e876..ec2e9f3d5e 100644 --- a/lightllm/common/basemodel/triton_kernel/linear_att/replayssm.py +++ b/lightllm/common/basemodel/triton_kernel/linear_att/replayssm.py @@ -291,7 +291,7 @@ def _replay( qk /= tl.sqrt(tl.sum(qk * qk, axis=1) + 1.0e-6)[:, None] else: qk *= tl.rsqrt(tl.sum(qk * qk, axis=1) + 1.0e-6)[:, None] - qk *= tl.where(tt < BW, K**-0.5, 1.0)[:, None] + qk *= tl.where(tt < BW, K ** -0.5, 1.0)[:, None] if KDA: history = ll < n gates = tl.load( @@ -474,7 +474,7 @@ def _replay( 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) + 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) if KDA: x = tl.load(A + t * SA + hv * K + kk, kk < K, 0).to(tl.float32) + bias diff --git a/lightllm/common/basemodel/triton_kernel/linear_att/replayssm_compact.py b/lightllm/common/basemodel/triton_kernel/linear_att/replayssm_compact.py index 9ca2b75ef6..113bb4e152 100644 --- a/lightllm/common/basemodel/triton_kernel/linear_att/replayssm_compact.py +++ b/lightllm/common/basemodel/triton_kernel/linear_att/replayssm_compact.py @@ -108,7 +108,7 @@ def _compact( state += k[:, None] * d[None, :] if not COMMIT: q = tl.load(Q + t * SQ + h * K + kk, kk < K, 0).to(tl.float32) - q = q / tl.sqrt(tl.sum(q * q) + 1.0e-6) * (K**-0.5) + 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: diff --git a/lightllm/common/kv_cache_mem_manager/mem_manager.py b/lightllm/common/kv_cache_mem_manager/mem_manager.py index 2d6e6e0105..fd5d814fe9 100755 --- a/lightllm/common/kv_cache_mem_manager/mem_manager.py +++ b/lightllm/common/kv_cache_mem_manager/mem_manager.py @@ -99,7 +99,7 @@ def profile_size(self, mem_fraction): available_memory = get_available_gpu_memory(world_size) - get_total_gpu_memory() * (1 - mem_fraction) cell_size = self.get_cell_size() pd_kv_move_buffer_size = self.get_pd_kv_move_buffer_size() - available_memory_bytes = available_memory * 1024**3 - pd_kv_move_buffer_size + available_memory_bytes = available_memory * 1024 ** 3 - pd_kv_move_buffer_size self.size = int(available_memory_bytes / cell_size) if world_size > 1: tensor = torch.tensor(self.size, dtype=torch.int64, device=f"cuda:{get_current_device_id()}") diff --git a/lightllm/server/router/model_infer/mode_backend/chunked_prefill/impl.py b/lightllm/server/router/model_infer/mode_backend/chunked_prefill/impl.py index 8ad3393f5f..7cefa1b677 100644 --- a/lightllm/server/router/model_infer/mode_backend/chunked_prefill/impl.py +++ b/lightllm/server/router/model_infer/mode_backend/chunked_prefill/impl.py @@ -117,12 +117,7 @@ def prefill_normal( event_pack.notify_post_handle_and_wait_pre_post_handle() with torch.cuda.stream(g_infer_context.get_overlap_stream()): - ( - _, - next_token_ids_cpu, - next_token_logprobs_cpu, - next_token_ranks_cpu, - ) = self._sample_and_scatter_token( + (_, next_token_ids_cpu, next_token_logprobs_cpu, next_token_ranks_cpu,) = self._sample_and_scatter_token( logits=model_output.logits, b_req_idx=model_input.b_req_idx, b_mtp_index=model_input.b_mtp_index, @@ -167,12 +162,7 @@ def decode_normal( event_pack.notify_post_handle_and_wait_pre_post_handle() with torch.cuda.stream(g_infer_context.get_overlap_stream()): - ( - _, - next_token_ids_cpu, - next_token_logprobs_cpu, - next_token_ranks_cpu, - ) = self._sample_and_scatter_token( + (_, next_token_ids_cpu, next_token_logprobs_cpu, next_token_ranks_cpu,) = self._sample_and_scatter_token( logits=model_output.logits, b_req_idx=model_input.b_req_idx, b_mtp_index=model_input.b_mtp_index, From f1e8288dff53b25a17c9376e60cb21dec4b008f4 Mon Sep 17 00:00:00 2001 From: sufubao Date: Tue, 20 Oct 2026 12:17:47 +0800 Subject: [PATCH 30/40] fix: bound SSM tuning memory and prepare remote configs --- .../triton_kernel/linear_att/ssm_autotune.py | 22 +++++++ lightllm/common/triton_utils/autotuner.py | 20 +++++- lightllm/server/api_start.py | 12 ++-- .../linear_att/test_ssm_autotune_memory.py | 65 +++++++++++++++++++ unit_tests/common/test_autotuner.py | 58 +++++++++++++++++ unit_tests/server/test_hybrid_start_args.py | 46 +++++++++++++ 6 files changed, 216 insertions(+), 7 deletions(-) create mode 100644 unit_tests/common/basemodel/triton_kernel/linear_att/test_ssm_autotune_memory.py diff --git a/lightllm/common/basemodel/triton_kernel/linear_att/ssm_autotune.py b/lightllm/common/basemodel/triton_kernel/linear_att/ssm_autotune.py index 6546764134..9f8618cc7d 100644 --- a/lightllm/common/basemodel/triton_kernel/linear_att/ssm_autotune.py +++ b/lightllm/common/basemodel/triton_kernel/linear_att/ssm_autotune.py @@ -55,9 +55,31 @@ def run_key(cache, q, cu_seqlens): return tokens, cu_seqlens.numel() - 1 if cu_seqlens is not None else tokens +def scratch_size_bytes(cache, batch, sequences, inputs): + # All cache tensors are request-owned; retain every layer and the HOLD slot. + slots = cache.state.shape[1] + pool_bytes = sum(x.numel() * x.element_size() for x in vars(cache).values() if isinstance(x, torch.Tensor)) + input_bytes = sum( + (1 + sum((size - 1) * stride for size, stride in zip(x.shape, x.stride()))) * x.element_size() + for x in inputs + if x.numel() + ) + snapshot_bytes = cache.state.numel() * cache.state.element_size() // slots + return pool_bytes // slots * (batch + 1) + input_bytes + snapshot_bytes + 32 * (batch + sequences + 1) + + def rebuild_inputs(cache, mode, q, k, v, a, b, a_log, bias, cu_seqlens=None, workload=None, run_config=None): tokens, sequences = run_key(cache, q, cu_seqlens) batch = min(sequences, triton.cdiv(tokens, cache.verify_width)) if cu_seqlens is not None else tokens + if q.is_cuda: + required = scratch_size_bytes(cache, batch, sequences, (q, k, v, a, b, a_log, bias)) + free, _ = torch.cuda.mem_get_info(q.device) + available = free + torch.cuda.memory_reserved(q.device) - torch.cuda.memory_allocated(q.device) + # Leave half for benchmark intermediates, CUDA Graphs and allocator fragmentation. + if required > available // 2: + raise torch.cuda.OutOfMemoryError( + f"SSM tuning scratch needs {required} bytes, exceeding its {available // 2}-byte budget" + ) layers, _, hv, kd, vd = cache.state.shape state = torch.empty((layers, batch + 1, hv, kd, vd), device=q.device, dtype=cache.state.dtype).normal_(0, 0.02) if mode == "replay": diff --git a/lightllm/common/triton_utils/autotuner.py b/lightllm/common/triton_utils/autotuner.py index b7a94d657e..c664bbe2b0 100644 --- a/lightllm/common/triton_utils/autotuner.py +++ b/lightllm/common/triton_utils/autotuner.py @@ -181,6 +181,7 @@ def __init__( self.cached_configs = {} self.fast_match_configs = collections.defaultdict(dict) self.warmuped_configs_set = set() + self._oom_keys = set() self.arg_names = [param.name for param in inspect.signature(self.fn).parameters.values()] self._argname_to_pos = {name: idx for idx, name in enumerate(self.arg_names)} self._pos_to_argname = {idx: name for idx, name in enumerate(self.arg_names)} @@ -244,6 +245,8 @@ def __call__(self, *args, **kwargs): need_tuning = (autotune_level == AutotuneLevel.FORCE_AUTOTUNE) or ( run_key not in self.cached_configs.get(static_key, {}) ) + if (static_key, run_key) in self._oom_keys: + need_tuning = False if world_size > 1: _need_tunings = [None for _ in range(world_size)] dist.all_gather_object(_need_tunings, obj=need_tuning, group=self._get_autotune_group()) @@ -390,7 +393,22 @@ def _autotune(self, args, kwargs, static_key, run_key, rank_id, world_size): # 仅为本次调优重建输入，构造开销不计入计时，最终执行仍使用调用方的原始输入。 if self.rebuild_input_func is not None: - args, kwargs = self.rebuild_input_func(*args, **kwargs) + ready = True + try: + args, kwargs = self.rebuild_input_func(*args, **kwargs) + except torch.cuda.OutOfMemoryError as error: + ready = False + logger.warning(f"Cannot allocate tuning inputs for {self.kernel_name}: {error}") + if world_size > 1: + all_ready = [None for _ in range(world_size)] + dist.all_gather_object(all_ready, obj=ready, group=self._get_autotune_group()) + ready = all(all_ready) + if not ready: + # Every rank skips the search. Keep existing/default configs without + # persisting a memory-pressure fallback as a measured tuning result. + self._oom_keys.add((static_key, run_key)) + logger.warning(f"Skipping autotune for {self.kernel_name} / {run_key}: insufficient scratch memory") + return best_config = None best_time = float("inf") diff --git a/lightllm/server/api_start.py b/lightllm/server/api_start.py index c1f7e92178..85cd8e7e6c 100644 --- a/lightllm/server/api_start.py +++ b/lightllm/server/api_start.py @@ -120,6 +120,12 @@ def _launch_subprocesses(args: StartArgs): not args.disable_cudagraph or args.run_mode == "prefill" ), "--disable_cudagraph is only supported on Prefill nodes when --mtp_mode is enabled" + # Prepare remote configuration before any model-dependent defaults or validation. + if args.model_dir and "s3://" in args.model_dir: + from lightllm.utils.petrel_helper import s3_model_prepare + + s3_model_prepare(args.model_dir) + auto_set_max_req_total_len(args) auto_set_fused_shared_experts(args) set_unique_server_name(args) @@ -389,12 +395,6 @@ def _launch_subprocesses(args: StartArgs): args.cpu_cache_token_page_size % args.page_size == 0 ), "--cpu_cache_token_page_size must be divisible by --page_size" - # help to manage data stored on Ceph - if "s3://" in args.model_dir: - from lightllm.utils.petrel_helper import s3_model_prepare - - s3_model_prepare(args.model_dir) - # 如果args.eos_id 是 None, 从 config.json 中读取 eos_token_id 相关的信息，赋值给 args if args.eos_id is None: from lightllm.utils.config_utils import get_eos_token_ids diff --git a/unit_tests/common/basemodel/triton_kernel/linear_att/test_ssm_autotune_memory.py b/unit_tests/common/basemodel/triton_kernel/linear_att/test_ssm_autotune_memory.py new file mode 100644 index 0000000000..022c71227b --- /dev/null +++ b/unit_tests/common/basemodel/triton_kernel/linear_att/test_ssm_autotune_memory.py @@ -0,0 +1,65 @@ +import pytest +import torch + +from lightllm.common.basemodel.triton_kernel.linear_att import ssm_autotune as tuning +from lightllm.common.basemodel.triton_kernel.linear_att.replayssm import ReplaySSMCache + + +@pytest.mark.parametrize("width", [1, 5]) +def test_full_layer_scratch_exceeds_h200_headroom_without_allocating(width): + state = torch.empty((48, 257, 48, 128, 128), dtype=torch.bfloat16, device="meta") + cache = ReplaySSMCache(state, 16, width, num_key_heads=16) + tokens = 256 * width + inputs = [ + torch.empty(shape, dtype=torch.bfloat16, device="meta") + for shape in [(tokens, 16, 128), (tokens, 16, 128), (tokens, 48, 128), (tokens, 48), (tokens, 48), (48,), (48,)] + ] + required = tuning.scratch_size_bytes(cache, 256, 256, inputs) + assert 36 * 1024 ** 3 < required < 37 * 1024 ** 3 + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") +@pytest.mark.parametrize("width", [1, 5]) +def test_tuning_budget_falls_back_without_touching_serving_state(monkeypatch, tmp_path, width): + from lightllm.common.triton_utils import autotuner as autotuner_module + from lightllm.common.triton_utils.autotuner import Autotuner, AutotuneKernelType + + state = torch.zeros((2, 4, 2, 32, 32), dtype=torch.bfloat16, device="cuda") + cache = ReplaySSMCache(state, 8, width, num_key_heads=2) + for tensor in vars(cache).values(): + if isinstance(tensor, torch.Tensor): + tensor.zero_() + tokens = 3 * width + q = torch.ones((tokens, 2, 32), dtype=torch.bfloat16, device="cuda") + gates = torch.zeros((tokens, 2), dtype=torch.bfloat16, device="cuda") + a_log = torch.zeros(2, device="cuda") + bias = torch.zeros(2, device="cuda") + cu = torch.arange(4, device="cuda", dtype=torch.int32) * width if width > 1 else None + if cu is not None: + q = q.unsqueeze(0) + inputs = dict(cache=cache, mode="replay", q=q, k=q, v=q, a=gates, b=gates, a_log=a_log, bias=bias, cu_seqlens=cu) + before = {name: x.clone() for name, x in vars(cache).items() if isinstance(x, torch.Tensor)} + monkeypatch.setattr(torch.cuda, "mem_get_info", lambda device: (0, 1024 ** 3)) + monkeypatch.setattr(torch.cuda, "memory_reserved", lambda device: 0) + monkeypatch.setattr(torch.cuda, "memory_allocated", lambda device: 0) + tuner = tuning.select_config + monkeypatch.setattr(tuner, "_cache_dir", str(tmp_path), raising=False) + monkeypatch.setattr(tuner, "cached_configs", {}) + monkeypatch.setattr(tuner, "fast_match_configs", {}) + monkeypatch.setattr(tuner, "_oom_keys", set()) + monkeypatch.setattr(autotuner_module, "get_triton_autotune_level", lambda: 1) + monkeypatch.setattr( + tuner, "_bench", lambda *args, **kwargs: pytest.fail("Insufficient memory must skip benchmarks") + ) + with Autotuner.autotune_warmup(AutotuneKernelType.DECODE_ATTENTION): + config = tuning.configure_cache(**inputs) + assert tuning.configure_cache(**inputs) is config + assert config == {"BV": 32, "num_warps": 1} + assert not list(tmp_path.iterdir()) + for name, expected in before.items(): + torch.testing.assert_close(getattr(cache, name), expected, atol=0, rtol=0) + + reqs = torch.arange(3, device="cuda", dtype=torch.int32) + cache.prepare_decode(reqs, cu) + cache.forward(0, q, q, q, gates, gates, a_log, bias, reqs, cu, run_config=config) + torch.cuda.synchronize() diff --git a/unit_tests/common/test_autotuner.py b/unit_tests/common/test_autotuner.py index 547246cfe5..5eb12f9e57 100644 --- a/unit_tests/common/test_autotuner.py +++ b/unit_tests/common/test_autotuner.py @@ -484,3 +484,61 @@ def test_default_api_and_nested_phase_restore_after_exception(): assert not Autotuner.is_autotune_warmup() assert not Autotuner.is_kernel_autotune_warmup(AutotuneKernelType.GENERAL) assert not Autotuner.is_kernel_autotune_warmup(AutotuneKernelType.DECODE_ATTENTION) + + +@pytest.mark.parametrize("world_size,failed_rank", [(1, 0), (2, 0), (2, 1)]) +@pytest.mark.parametrize("cached", [False, True]) +def test_rebuild_oom_skips_search_on_all_ranks_without_saving_fallback( + tmp_path, monkeypatch, world_size, failed_rank, cached +): + kernel, calls, benchmarks, cache_file = make_kernel( + tmp_path, monkeypatch, "scratch_oom", AutotuneKernelType.DECODE_ATTENTION, warmup_all_exist_config=False + ) + if cached: + cache_file.write_text(json.dumps({"8": {"block": 1}})) + before = cache_file.read_bytes() if cached else None + rebuilds = [] + + def rebuild(size, run_config=None): + rebuilds.append(size) + if failed_rank == 0: + raise torch.cuda.OutOfMemoryError("scratch allocation") + return (size,), {} + + monkeypatch.setattr(kernel, "rebuild_input_func", rebuild) + monkeypatch.setattr(autotuner_module.dist, "is_initialized", lambda: world_size > 1) + monkeypatch.setattr(autotuner_module, "get_global_rank", lambda: 0) + monkeypatch.setattr(autotuner_module, "get_global_world_size", lambda: world_size) + monkeypatch.setattr(kernel, "_get_autotune_group", lambda: None) + collectives = [] + + def all_gather_object(outputs, obj, group): + collectives.append(obj) + if len(collectives) == 3 and isinstance(obj, bool): + outputs[:] = [rank != failed_rank for rank in range(world_size)] + else: + outputs[:] = [obj] * world_size + + monkeypatch.setattr(autotuner_module.dist, "all_gather_object", all_gather_object) + fallback = {"block": 1} if cached else None + with Autotuner.autotune_warmup(AutotuneKernelType.DECODE_ATTENTION): + assert kernel(16) == fallback + assert kernel(16) == fallback + assert rebuilds == [16] + assert calls == [(16, fallback), (16, fallback)] + assert not benchmarks + assert (cache_file.read_bytes() if cache_file.exists() else None) == before + assert "16" not in kernel.cached_configs.get(frozendict({"dtype": "test"}), {}) + if world_size > 1: + assert len(collectives) == 4 # need, key, allocation status, repeated need + + +def test_rebuild_programming_error_is_not_treated_as_oom(tmp_path, monkeypatch): + kernel, _, _, _ = make_kernel(tmp_path, monkeypatch, "rebuild_error") + + def rebuild(size, run_config=None): + raise ValueError("invalid scratch layout") + + monkeypatch.setattr(kernel, "rebuild_input_func", rebuild) + with Autotuner.autotune_warmup(), pytest.raises(ValueError, match="invalid scratch layout"): + kernel(16) diff --git a/unit_tests/server/test_hybrid_start_args.py b/unit_tests/server/test_hybrid_start_args.py index 230e23b0e5..4959a4ac65 100644 --- a/unit_tests/server/test_hybrid_start_args.py +++ b/unit_tests/server/test_hybrid_start_args.py @@ -102,3 +102,49 @@ def test_whole_image_blocks_fit_batch_at_startup( else: with pytest.raises(AssertionError, match="batch_max_tokens >= max_image_token_count"): _launch_subprocesses(args) + + +@pytest.mark.parametrize("mode,step", [("native", 0), ("compact", 4), ("replay", 0)]) +def test_s3_configuration_is_prepared_before_ssm_validation(launch_args, monkeypatch, tmp_path, mode, step): + from pathlib import Path + from lightllm.utils.petrel_helper import PetrelHelper + + monkeypatch.chdir(tmp_path) + model_dir = "s3://test-bucket/qwen" + config = {"model_type": "qwen3_5_text"} + downloads = [] + + def list_files(path, extension): + assert path == model_dir + return ["config.json"] if extension == "json" else [] + + def download_json(remote, local): + downloads.append(remote) + Path(local).write_text(json.dumps(config)) + + monkeypatch.setattr(PetrelHelper, "list", list_files) + monkeypatch.setattr(PetrelHelper, "download_json", download_json) + args = launch_args("qwen3_5_text", ssm_state_mode=mode, mtp_step=step, mtp_mode="eagle_with_att" if step else None) + args.model_dir = model_dir + assert not Path(model_dir, "config.json").exists() + with pytest.raises(RuntimeError, match="startup validation passed"): + _launch_subprocesses(args) + assert downloads == [model_dir + "/config.json"] + if mode == "replay": + assert args.ssm_replay_config == {"cache_len": 16, "compute_mode": "inline", "fold_method": "raw"} + + +def test_cold_s3_model_still_enforces_hybrid_constraints(launch_args, monkeypatch, tmp_path): + from pathlib import Path + + monkeypatch.chdir(tmp_path) + + def prepare(model_dir): + Path(model_dir).mkdir(parents=True) + Path(model_dir, "config.json").write_text(json.dumps({"model_type": "qwen3_5_text"})) + + monkeypatch.setattr("lightllm.utils.petrel_helper.s3_model_prepare", prepare) + args = launch_args("qwen3_5_text", disable_dynamic_prompt_cache=True) + args.model_dir = "s3://test-bucket/qwen" + with pytest.raises(AssertionError, match="hybrid attention models require dynamic prompt cache"): + _launch_subprocesses(args) From 924088d080a6e6472bbdf94996d26c3fbf25b6b4 Mon Sep 17 00:00:00 2001 From: sufubao Date: Wed, 21 Oct 2026 18:16:47 +0800 Subject: [PATCH 31/40] refactor: simplify ReplaySSM state interfaces --- .../basemodel/attention/linear/flashqla.py | 4 +- .../common/basemodel/attention/linear/gdn.py | 150 ++++------ .../common/basemodel/attention/linear/kda.py | 85 ++---- .../basemodel/attention/linear/triton.py | 4 +- lightllm/common/basemodel/basemodel.py | 30 +- lightllm/common/basemodel/batch_objs.py | 10 +- lightllm/common/basemodel/cuda_graph.py | 8 +- lightllm/common/basemodel/infer_struct.py | 6 +- .../{replayssm_compact.py => compact_ssm.py} | 59 ++-- .../linear_att/mtp_state_params.py | 28 +- .../triton_kernel/linear_att/replayssm.py | 90 +++--- .../triton_kernel/linear_att/ssm_autotune.py | 133 ++++---- .../triton_kernel/linear_att_copy.py | 14 +- .../post_process/apply_penalty_gpu_cache.py | 41 ++- .../triton_kernel/post_process/argmax.py | 47 --- .../triton_kernel/post_process/greedy.py | 66 ---- .../qwen3next_mem_manager.py | 10 +- .../{recurrent_state.py => linear_state.py} | 166 +++++----- lightllm/common/recurrent_state_types.py | 81 ----- lightllm/common/req_manager/hybrid_base.py | 2 +- lightllm/common/req_manager/linear_att.py | 20 +- .../llama/layer_infer/post_layer_infer.py | 4 +- lightllm/server/api_start.py | 4 +- .../model_infer/mode_backend/base_backend.py | 16 +- .../mode_backend/chunked_prefill/impl.py | 2 +- .../mode_backend/dp_backend/impl.py | 4 +- .../mode_backend/generic_post_process.py | 23 +- .../proposers/eagle_with_att.py | 8 +- .../model_infer/mtp_speculative/utils.py | 4 +- test/benchmark/kernels/benchmark_kda_ssm.py | 54 ++-- .../kernels/benchmark_replayssm_service.py | 10 +- .../common/basemodel/test_draft_tail_graph.py | 18 +- .../common/basemodel/test_model_output.py | 6 +- .../linear_att/test_acceptance_fusion.py | 136 +++++---- .../linear_att/test_mtp_state_params.py | 16 +- .../linear_att/test_replayssm.py | 283 +++++++++--------- .../linear_att/test_replayssm_execution.py | 154 +++++----- .../linear_att/test_replayssm_grouped_keys.py | 48 +-- .../linear_att/test_replayssm_parallel.py | 81 ++--- .../linear_att/test_ssm_autotune.py | 226 +++++++------- .../linear_att/test_ssm_autotune_memory.py | 32 +- .../triton_kernel/test_greedy_sampling.py | 143 --------- ...ecurrent_state.py => test_linear_state.py} | 48 +-- ...tract.py => test_linear_state_contract.py} | 45 ++- unit_tests/models/test_post_layer_output.py | 4 +- .../mtp_speculative/test_eagle_with_att.py | 16 +- unit_tests/server/test_ssm_state_mode.py | 34 +-- 47 files changed, 1072 insertions(+), 1401 deletions(-) rename lightllm/common/basemodel/triton_kernel/linear_att/{replayssm_compact.py => compact_ssm.py} (85%) delete mode 100644 lightllm/common/basemodel/triton_kernel/post_process/argmax.py delete mode 100644 lightllm/common/basemodel/triton_kernel/post_process/greedy.py rename lightllm/common/{recurrent_state.py => linear_state.py} (57%) delete mode 100644 lightllm/common/recurrent_state_types.py delete mode 100644 unit_tests/common/basemodel/triton_kernel/test_greedy_sampling.py rename unit_tests/common/{test_recurrent_state.py => test_linear_state.py} (77%) rename unit_tests/common/{test_recurrent_state_contract.py => test_linear_state_contract.py} (66%) diff --git a/lightllm/common/basemodel/attention/linear/flashqla.py b/lightllm/common/basemodel/attention/linear/flashqla.py index 6e57eca732..79ad4e20e6 100644 --- a/lightllm/common/basemodel/attention/linear/flashqla.py +++ b/lightllm/common/basemodel/attention/linear/flashqla.py @@ -1,7 +1,7 @@ -from lightllm.common.basemodel.attention.linear.gdn import LinearAttBackend +from lightllm.common.basemodel.attention.linear.gdn import GDNLinearAttBackend -class FlashQlaLinearAttBackend(LinearAttBackend): +class FlashQlaLinearAttBackend(GDNLinearAttBackend): def get_prefill_kernel(self): from flash_qla import chunk_gated_delta_rule diff --git a/lightllm/common/basemodel/attention/linear/gdn.py b/lightllm/common/basemodel/attention/linear/gdn.py index e02904aac3..cdd6016a31 100644 --- a/lightllm/common/basemodel/attention/linear/gdn.py +++ b/lightllm/common/basemodel/attention/linear/gdn.py @@ -6,9 +6,7 @@ import torch -from lightllm.common.recurrent_state_types import SSMRunConfig from ..base_att import BaseAttBackend, BasePrefillAttState, BaseDecodeAttState, AttControl -from lightllm.common.req_manager import ReqManagerForMamba from lightllm.utils.envs_utils import get_env_start_args from lightllm.common.basemodel.triton_kernel.linear_att.causal_conv1d import causal_conv1d_fn from lightllm.common.basemodel.triton_kernel.linear_att.fused_gdn_gating import fused_gdn_gating @@ -30,13 +28,9 @@ from lightllm.models.qwen3next.layer_infer.transformer_layer_infer import Qwen3NextTransformerLayerWeight -class LinearAttBackend(BaseAttBackend, ABC): +class GDNLinearAttBackend(BaseAttBackend, ABC): def __init__(self, model: "TpPartBaseModel"): super().__init__(model=model) - req_manager = model.req_manager - if not isinstance(req_manager, ReqManagerForMamba): - raise TypeError(f"LinearAttBackend requires ReqManagerForMamba, got {type(req_manager).__name__}") - self.req_manager: ReqManagerForMamba = req_manager self._init_linear_layer_metadata(network_config=model.config, tp_world_size=model.tp_world_size_) self.prefill_kernel = self.get_prefill_kernel() @@ -107,16 +101,16 @@ def _rearrange_mixed_qkv(self, mixed_qkv, decode=False): value = value.view(1, seq_len, self.tp_num_v_heads, self.head_v_dim) return query, key, value - def create_att_prefill_state(self, infer_state: "InferStateInfo") -> "LinearAttPrefillAttState": - return LinearAttPrefillAttState(backend=self, _infer_state_weak_ref=weakref.ref(infer_state)) + def create_att_prefill_state(self, infer_state: "InferStateInfo") -> "GDNLinearAttPrefillAttState": + return GDNLinearAttPrefillAttState(backend=self, _infer_state_weak_ref=weakref.ref(infer_state)) - def create_att_decode_state(self, infer_state: "InferStateInfo") -> "LinearAttDecodeAttState": - return LinearAttDecodeAttState(backend=self, _infer_state_weak_ref=weakref.ref(infer_state)) + def create_att_decode_state(self, infer_state: "InferStateInfo") -> "GDNLinearAttDecodeAttState": + return GDNLinearAttDecodeAttState(backend=self, _infer_state_weak_ref=weakref.ref(infer_state)) @dataclasses.dataclass -class LinearAttPrefillAttState(BasePrefillAttState): - backend: LinearAttBackend = None +class GDNLinearAttPrefillAttState(BasePrefillAttState): + backend: GDNLinearAttBackend = None b_conv_buffer_idx: torch.Tensor = None b_ssm_buffer_idx: torch.Tensor = None @@ -124,8 +118,8 @@ def init_state(self): # 每次 _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.backend.req_manager.recurrent_state.prepare_prefill(self.infer_state.b_req_idx) + linear_state_manager = self.infer_state.req_manager.linear_state_manager + self.b_conv_buffer_idx, self.b_ssm_buffer_idx = linear_state_manager.prepare_prefill(self.infer_state.b_req_idx) return def prefill_att( @@ -143,9 +137,9 @@ def prefill_att( mixed_qkvzba: torch.Tensor = linear_att_dict["mixed_qkvzba"] layer_weight = linear_att_dict["layer_weight"] layer_num = linear_att_dict["layer_num"] - backend: LinearAttBackend = self.backend + backend: GDNLinearAttBackend = self.backend - conv_states, ssm_states = backend.req_manager.get_mamba_cache(layer_num) + conv_states, ssm_states = self.infer_state.req_manager.get_mamba_cache(layer_num) # 在开启了mtp的时候，conv 状态的最后一维可能存在冗余的部分，需要进行切片对齐。 # prefill 模式下，使用不到这几个维度，所以需要扣除掉， if backend.mtp_step > 0: @@ -169,7 +163,7 @@ def _gdn_prefill_kernel( g, beta = fused_gdn_gating(layer_weight.linear_A_log.weight, a, b, layer_weight.linear_dt_bias.weight) mixed_qkv = mixed_qkv.transpose(0, 1) - backend: LinearAttBackend = self.backend + backend: GDNLinearAttBackend = self.backend out_tensor = causal_conv1d_fn( mixed_qkv, layer_weight.linear_conv1d.mm_param.weight, @@ -205,8 +199,8 @@ def _gdn_prefill_kernel( @dataclasses.dataclass -class LinearAttDecodeAttState(BaseDecodeAttState): - backend: LinearAttBackend = None +class GDNLinearAttDecodeAttState(BaseDecodeAttState): + backend: GDNLinearAttBackend = None b_conv_buffer_idx: torch.Tensor = None b_ssm_buffer_idx: torch.Tensor = None b1_mtp_cu_q_seq_len: torch.Tensor = None @@ -216,41 +210,30 @@ 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: self._init_normal_decode_state() - return - - # MTP 元数据按模型初始化，沿用每个角色只有一个 GDN State 的约束。 - # 只在此读取游标；Backend 不保存 State，逐层计算按逻辑层号查询资源。 - (state,) = ( - state - for state in self.infer_state.att_state_manager.get_bounded_attention_states( - self.infer_state.att_state_role - ) - if isinstance(state, GDNState) - ) - if self.backend.uses_dynamic_spec_verify_layout(): - self._init_dynamic_mtp_decode_state(draft_step + 1, req_to_mtp_state_index=state.mtp_state_index) + elif self.backend.uses_dynamic_spec_verify_layout(): + self._init_dynamic_mtp_decode_state() else: self._init_fixed_mtp_decode_state(draft_step) - self.backend.req_manager.recurrent_state.prepare_decode( - self.b_conv_buffer_idx, self.b1_mtp_cu_q_seq_len, single_token=draft_step == 0 - ) + linear_state_manager = self.infer_state.req_manager.linear_state_manager + if draft_step == 0: + linear_state_manager.prepare_decode(self.b_conv_buffer_idx, single_token=True) + else: + ( + self.b_conv_buffer_idx, + self.b_ssm_buffer_idx, + self.b_num_accepted_tokens, + ) = linear_state_manager.prepare_mtp_decode(self.b_conv_buffer_idx, self.b1_mtp_cu_q_seq_len) def _init_normal_decode_state(self): self.b_conv_buffer_idx = self.infer_state.b_req_idx self.b_ssm_buffer_idx = self.infer_state.b_req_idx - def _init_dynamic_mtp_decode_state(self, mtp_size: int, *, req_to_mtp_state_index: torch.Tensor) -> None: - ( - self.b1_mtp_cu_q_seq_len, - self.b_conv_buffer_idx, - self.b_num_accepted_tokens, - ) = build_dynamic_mtp_linear_att_state_params( + def _init_dynamic_mtp_decode_state(self): + (self.b1_mtp_cu_q_seq_len, self.b_conv_buffer_idx,) = build_dynamic_mtp_linear_att_state_params( b_req_idx=self.infer_state.b_req_idx, b_mtp_index=self.infer_state.b_mtp_index, - req_to_mtp_state_index=self.backend.req_manager.recurrent_state.accepted_index, - hold_req_id=self.backend.req_manager.HOLD_REQUEST_ID, + hold_req_id=self.infer_state.req_manager.HOLD_REQUEST_ID, ) - self._init_mtp_ssm_buffer_idx(mtp_size) def _init_fixed_mtp_decode_state(self, draft_step: int, *, req_to_mtp_state_index: torch.Tensor) -> None: mtp_size = draft_step + 1 @@ -269,22 +252,6 @@ def _init_fixed_mtp_decode_state(self, draft_step: int, *, req_to_mtp_state_inde device=self.infer_state.b_req_idx.device, ) self.b_conv_buffer_idx = self.infer_state.b_req_idx.view(att_batch_size, mtp_size)[:, 0].contiguous() - self.b_num_accepted_tokens = self.backend.req_manager.recurrent_state.accepted_index[self.b_conv_buffer_idx] + 1 - self._init_mtp_ssm_buffer_idx(mtp_size) - - def _init_mtp_ssm_buffer_idx(self, mtp_size: int): - if self.backend.req_manager.recurrent_state.update_backend 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) - state_offsets = torch.arange( - mtp_size, - device=self.infer_state.b_req_idx.device, - dtype=self.infer_state.b_req_idx.dtype, - ).view(1, mtp_size) - self.b_ssm_buffer_idx = b_ssm_buffer_start_idx + state_offsets # [att_batch_size, mtp_size] def decode_att( self, @@ -301,10 +268,10 @@ def decode_att( mixed_qkvzba = linear_att_dict["mixed_qkvzba"] layer_weight = linear_att_dict["layer_weight"] layer_num = linear_att_dict["layer_num"] - backend: LinearAttBackend = self.backend + backend: GDNLinearAttBackend = self.backend mixed_qkv, z, b, a = backend._split_qkvzba(mixed_qkvzba) - conv_states, ssm_states = backend.req_manager.get_mamba_cache(layer_num) + conv_states, ssm_states = self.infer_state.req_manager.get_mamba_cache(layer_num) draft_step = self.backend.model.mtp_manager.get_decode_draft_step(self.backend.model.is_mtp_draft_model) if draft_step > 0: @@ -342,7 +309,7 @@ def _gdn_decode_kernel( infer_state: "Qwen3NextInferStateInfo", layer_weight: "Qwen3NextTransformerLayerWeight", ): - backend: LinearAttBackend = self.backend + backend: GDNLinearAttBackend = self.backend # Recurrent processing with fused gating. Decode uses a specialized # conv+pack kernel to avoid materializing the post-conv qkv tensor @@ -363,21 +330,24 @@ def _gdn_decode_kernel( backend.tp_num_v_heads, backend.head_v_dim, ) - updater = backend.req_manager.recurrent_state.update_backend - if updater is not None: - linear_layer_idx = backend.req_manager.linear_config.get_linear_layer_index(layer_weight.layer_num_) - args = (query, key, value, a, b, layer_weight.linear_A_log.weight, layer_weight.linear_dt_bias.weight) - if linear_layer_idx == 0: - self.infer_state.ssm_run_config = SSMRunConfig.from_mapping(updater.get_run_config(*args)) - return ( - updater.forward( - linear_layer_idx, - *args, - self.b_conv_buffer_idx, - run_config=self.infer_state.ssm_run_config, - ), - z, + ssm_backend = self.infer_state.req_manager.linear_state_manager.ssm_backend + if ssm_backend is not None: + linear_layer_idx = self.infer_state.req_manager.linear_config.get_linear_layer_index( + layer_weight.layer_num_ + ) + core_attn_out, self.infer_state.ssm_run_config = ssm_backend.forward( + linear_layer_idx, + query, + key, + value, + a, + b, + layer_weight.linear_A_log.weight, + layer_weight.linear_dt_bias.weight, + self.b_conv_buffer_idx, + run_config=self.infer_state.ssm_run_config, ) + return core_attn_out, z core_attn_out, _ = fused_recurrent_gated_delta_rule( q=query, k=key, @@ -407,7 +377,7 @@ def _gdn_mtp_kernel( causal_conv1d_update as causal_conv1d_update_mtp, ) - backend: LinearAttBackend = self.backend + backend: GDNLinearAttBackend = self.backend cu_seqlens_q = self.b1_mtp_cu_q_seq_len mixed_qkv = causal_conv1d_update_mtp( @@ -423,19 +393,25 @@ def _gdn_mtp_kernel( ) query, key, value = backend._rearrange_mixed_qkv(mixed_qkv, decode=False) - updater = backend.req_manager.recurrent_state.update_backend - if updater is not None: - linear_layer_idx = backend.req_manager.linear_config.get_linear_layer_index(layer_weight.layer_num_) - args = (query, key, value, a, b, layer_weight.linear_A_log.weight, layer_weight.linear_dt_bias.weight) - if linear_layer_idx == 0: - self.infer_state.ssm_run_config = SSMRunConfig.from_mapping(updater.get_run_config(*args, cu_seqlens_q)) - return updater.forward( + ssm_backend = self.infer_state.req_manager.linear_state_manager.ssm_backend + if ssm_backend is not None: + linear_layer_idx = self.infer_state.req_manager.linear_config.get_linear_layer_index( + layer_weight.layer_num_ + ) + core_attn_out, self.infer_state.ssm_run_config = ssm_backend.forward( linear_layer_idx, - *args, + query, + key, + value, + a, + b, + layer_weight.linear_A_log.weight, + layer_weight.linear_dt_bias.weight, self.b_conv_buffer_idx, cu_seqlens_q, run_config=self.infer_state.ssm_run_config, ) + return core_attn_out 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/attention/linear/kda.py b/lightllm/common/basemodel/attention/linear/kda.py index 7bccdd3126..e28058826b 100644 --- a/lightllm/common/basemodel/attention/linear/kda.py +++ b/lightllm/common/basemodel/attention/linear/kda.py @@ -10,7 +10,6 @@ import torch -from lightllm.common.recurrent_state_types import SSMRunConfig from lightllm.common.attention_state.bounded_att import KDAState from lightllm.common.basemodel.attention.base_att import ( AttControl, @@ -50,24 +49,23 @@ def __init__(self, model: "TpPartBaseModel"): self.mtp_step = model.args.mtp_step def create_att_prefill_state(self, infer_state: "InferStateInfo"): - return KDAPrefillAttState(backend=self, _infer_state_weak_ref=weakref.ref(infer_state)) + return KDALinearAttPrefillAttState(backend=self, _infer_state_weak_ref=weakref.ref(infer_state)) def create_att_decode_state(self, infer_state: "InferStateInfo"): - return KDADecodeAttState(backend=self, _infer_state_weak_ref=weakref.ref(infer_state)) + return KDALinearAttDecodeAttState(backend=self, _infer_state_weak_ref=weakref.ref(infer_state)) def split_qkv(self, mixed_qkv: torch.Tensor): return mixed_qkv.split(self.tp_hidden_size, dim=-1) @dataclasses.dataclass -class KDAPrefillAttState(BasePrefillAttState): +class KDALinearAttPrefillAttState(BasePrefillAttState): b_conv_buffer_idx: torch.Tensor = None b_ssm_buffer_idx: torch.Tensor = None def init_state(self): - self.b_conv_buffer_idx = self.infer_state.b_req_idx - manager = self.infer_state.req_manager - self.b_ssm_buffer_idx = manager.recurrent_state.prepare_prefill(self.infer_state.b_req_idx) + linear_state_manager = self.infer_state.req_manager.linear_state_manager + self.b_conv_buffer_idx, self.b_ssm_buffer_idx = linear_state_manager.prepare_prefill(self.infer_state.b_req_idx) def prefill_att( self, @@ -133,7 +131,7 @@ def prefill_att( @dataclasses.dataclass -class KDADecodeAttState(BaseDecodeAttState): +class KDALinearAttDecodeAttState(BaseDecodeAttState): b_conv_buffer_idx: torch.Tensor = None b_ssm_buffer_idx: torch.Tensor = None b1_mtp_cu_q_seq_len: torch.Tensor = None @@ -143,41 +141,30 @@ def init_state(self): mtp_step = self.backend.mtp_step if mtp_step == 0: self._init_normal_decode_state() - return - - # MTP 元数据按模型初始化，沿用每个角色只有一个 KDA State 的约束。 - # 只在此读取游标；Backend 不保存 State，逐层计算按逻辑层号查询资源。 - (state,) = ( - state - for state in self.infer_state.att_state_manager.get_bounded_attention_states( - self.infer_state.att_state_role - ) - if isinstance(state, KDAState) - ) - if self.backend.uses_dynamic_spec_verify_layout(): - self._init_dynamic_mtp_decode_state(mtp_step + 1, req_to_mtp_state_index=state.mtp_state_index) + elif self.backend.uses_dynamic_spec_verify_layout(): + self._init_dynamic_mtp_decode_state() else: self._init_fixed_mtp_decode_state(mtp_step) - self.infer_state.req_manager.recurrent_state.prepare_decode( - self.b_conv_buffer_idx, self.b1_mtp_cu_q_seq_len, single_token=mtp_step == 0 - ) + linear_state_manager = self.infer_state.req_manager.linear_state_manager + if mtp_step == 0: + linear_state_manager.prepare_decode(self.b_conv_buffer_idx, single_token=True) + else: + ( + self.b_conv_buffer_idx, + self.b_ssm_buffer_idx, + self.b_num_accepted_tokens, + ) = linear_state_manager.prepare_mtp_decode(self.b_conv_buffer_idx, self.b1_mtp_cu_q_seq_len) def _init_normal_decode_state(self): self.b_conv_buffer_idx = self.infer_state.b_req_idx self.b_ssm_buffer_idx = self.infer_state.b_req_idx - def _init_dynamic_mtp_decode_state(self, mtp_size: int, *, req_to_mtp_state_index: torch.Tensor) -> None: - ( - self.b1_mtp_cu_q_seq_len, - self.b_conv_buffer_idx, - self.b_num_accepted_tokens, - ) = build_dynamic_mtp_linear_att_state_params( + def _init_dynamic_mtp_decode_state(self): + (self.b1_mtp_cu_q_seq_len, self.b_conv_buffer_idx,) = build_dynamic_mtp_linear_att_state_params( b_req_idx=self.infer_state.b_req_idx, b_mtp_index=self.infer_state.b_mtp_index, - req_to_mtp_state_index=self.infer_state.req_manager.recurrent_state.accepted_index, hold_req_id=self.infer_state.req_manager.HOLD_REQUEST_ID, ) - self._init_mtp_ssm_buffer_idx(mtp_size) def _init_fixed_mtp_decode_state(self, mtp_step: int, *, req_to_mtp_state_index: torch.Tensor) -> None: mtp_size = mtp_step + 1 @@ -196,24 +183,6 @@ def _init_fixed_mtp_decode_state(self, mtp_step: int, *, req_to_mtp_state_index: device=self.infer_state.b_req_idx.device, ) self.b_conv_buffer_idx = self.infer_state.b_req_idx.view(att_batch_size, mtp_size)[:, 0].contiguous() - self.b_num_accepted_tokens = ( - self.infer_state.req_manager.recurrent_state.accepted_index[self.b_conv_buffer_idx] + 1 - ) - self._init_mtp_ssm_buffer_idx(mtp_size) - - def _init_mtp_ssm_buffer_idx(self, mtp_size: int): - if self.infer_state.req_manager.recurrent_state.update_backend 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) - state_offsets = torch.arange( - mtp_size, - device=self.infer_state.b_req_idx.device, - dtype=self.infer_state.b_req_idx.dtype, - ).view(1, mtp_size) - self.b_ssm_buffer_idx = b_ssm_buffer_start_idx + state_offsets # [att_batch_size, mtp_size] def decode_att( self, @@ -256,10 +225,11 @@ def decode_att( v = v.view(*shape, backend.tp_num_heads, backend.head_dim) raw_gate = raw_gate.view(*shape, backend.tp_hidden_size) raw_beta = raw_beta.view(*shape, backend.tp_num_heads) - updater = self.infer_state.req_manager.recurrent_state.update_backend - if updater is not None: + ssm_backend = self.infer_state.req_manager.linear_state_manager.ssm_backend + if ssm_backend is not None: linear_layer_idx = self.infer_state.req_manager.linear_config.get_linear_layer_index(layer_num) - args = ( + output, self.infer_state.ssm_run_config = ssm_backend.forward( + linear_layer_idx, q, k, v, @@ -267,18 +237,11 @@ def decode_att( raw_beta.view(-1, backend.tp_num_heads), layer_weight.linear_A_log.weight, layer_weight.linear_dt_bias.weight, - ) - if linear_layer_idx == 0: - self.infer_state.ssm_run_config = SSMRunConfig.from_mapping( - updater.get_run_config(*args, self.b1_mtp_cu_q_seq_len) - ) - return updater.forward( - linear_layer_idx, - *args, self.b_conv_buffer_idx, self.b1_mtp_cu_q_seq_len, run_config=self.infer_state.ssm_run_config, ) + return output output, _ = fused_recurrent_kda( q=q, k=k, diff --git a/lightllm/common/basemodel/attention/linear/triton.py b/lightllm/common/basemodel/attention/linear/triton.py index f0b6400f51..f579f33a6c 100644 --- a/lightllm/common/basemodel/attention/linear/triton.py +++ b/lightllm/common/basemodel/attention/linear/triton.py @@ -1,7 +1,7 @@ -from lightllm.common.basemodel.attention.linear.gdn import LinearAttBackend +from lightllm.common.basemodel.attention.linear.gdn import GDNLinearAttBackend from lightllm.common.basemodel.triton_kernel.linear_att.fla.ops import chunk_gated_delta_rule -class TritonLinearAttBackend(LinearAttBackend): +class TritonLinearAttBackend(GDNLinearAttBackend): def get_prefill_kernel(self): return chunk_gated_delta_rule diff --git a/lightllm/common/basemodel/basemodel.py b/lightllm/common/basemodel/basemodel.py index f92d460f66..ddf7486c2e 100755 --- a/lightllm/common/basemodel/basemodel.py +++ b/lightllm/common/basemodel/basemodel.py @@ -360,7 +360,7 @@ def _align_decode_batch_size(self, batch_size: int) -> int: alignment = math.lcm(alignment, self.tp_world_size_) return triton.cdiv(batch_size, alignment) * alignment - def use_replayssm_draft_tail(self): + def should_project_draft_tail_logits(self): """Project only accepted-tail rows in the first draft pass.""" return ( self.is_mtp_draft_model @@ -397,7 +397,7 @@ def _init_cudagraph(self): ) self.draft_tail_graph = None if self.graph is not None: - if self.use_replayssm_draft_tail(): + if self.should_project_draft_tail_logits(): width = self.args.mtp_step + 1 self.draft_tail_graph = CudaGraph( batch_step_size_before_split=width, @@ -409,7 +409,7 @@ def _init_cudagraph(self): ) # Draft passes execute sequentially, so both layouts can share the pool. self.draft_tail_graph.mempool = self.graph.mempool - self.draft_tail_graph.warmup(self, logits_row_stride=width) + self.draft_tail_graph.warmup(self, input_rows_per_logit=width) if get_env_start_args().enable_decode_microbatch_overlap: self.graph.warmup_overlap(self) else: @@ -468,16 +468,16 @@ def _create_inferstate(self, model_input: ModelInput, microbatch_index: int = 0) infer_state.att_state_role = self.att_state_role infer_state.return_all_prompt_logics = self.return_all_prompt_logics infer_state.batch_size = model_input.batch_size - if model_input.b_logits_rows is not None: - assert self.use_replayssm_draft_tail() and not model_input.is_prefill + if model_input.b_logits_row_idx is not None: + assert self.should_project_draft_tail_logits() and not model_input.is_prefill width = self.args.mtp_step + 1 assert model_input.batch_size % width == 0 count = model_input.batch_size // width - rows = model_input.b_logits_rows + rows = model_input.b_logits_row_idx assert rows.ndim == 1 and rows.numel() <= count rows = F.pad(rows, (0, count - rows.numel())) - infer_state.b_logits_rows = rows - infer_state.logits_row_stride = width + infer_state.b_logits_row_idx = rows + infer_state.input_rows_per_logit = width infer_state.total_token_num = model_input.total_token_num infer_state.max_q_seq_len = model_input.max_q_seq_len infer_state.max_kv_seq_len = model_input.max_kv_seq_len @@ -624,15 +624,15 @@ def _create_model_output(self, post_output: PostLayerOutput, infer_state: InferS prompt_logics=infer_state.prompt_logics, logits_token_ids=post_output.logits_token_ids, ssm_run_config=infer_state.ssm_run_config, - logits_row_stride=infer_state.logits_row_stride, + input_rows_per_logit=infer_state.input_rows_per_logit, ) return output def _create_unpad_decode_model_output(self, model_output: ModelOutput, origin_batch_size: int): - stride = model_output.logits_row_stride - padded_batch_size = model_output.logits.shape[0] * stride - assert origin_batch_size % stride == 0 - logits_size = origin_batch_size // stride + input_rows_per_logit = model_output.input_rows_per_logit + padded_batch_size = model_output.logits.shape[0] * input_rows_per_logit + assert origin_batch_size % input_rows_per_logit == 0 + logits_size = origin_batch_size // input_rows_per_logit if padded_batch_size == origin_batch_size: return model_output new_model_output = copy.copy(model_output) @@ -732,10 +732,10 @@ def _decode( ) origin_batch_size = model_input.batch_size - graph = self.draft_tail_graph if model_input.b_logits_rows is not None else self.graph + graph = self.draft_tail_graph if model_input.b_logits_row_idx is not None else self.graph # 空 DP rank 也需要完整的 dummy verify group；TP/SP 切分不能破坏 MTP 分组。 infer_batch_size = self._align_decode_batch_size(max(1, origin_batch_size)) - if model_input.b_logits_rows is not None: + if model_input.b_logits_row_idx is not None: width = self.args.mtp_step + 1 infer_batch_size = triton.cdiv(infer_batch_size, width) * width diff --git a/lightllm/common/basemodel/batch_objs.py b/lightllm/common/basemodel/batch_objs.py index db678103e1..d2da5d377f 100644 --- a/lightllm/common/basemodel/batch_objs.py +++ b/lightllm/common/basemodel/batch_objs.py @@ -7,7 +7,7 @@ from lightllm.utils.tensor_utils import tensor_to_no_ref_tensor if TYPE_CHECKING: - from lightllm.common.recurrent_state_types import SSMRunConfig + from lightllm.common.linear_state import SSMRunConfig @dataclass @@ -56,12 +56,12 @@ class ModelInput: mtp_draft_input_hiddens: Optional[torch.Tensor] = None # Optional accepted-tail rows for the first EAGLE draft projection. - b_logits_rows: Optional[torch.Tensor] = None + b_logits_row_idx: Optional[torch.Tensor] = None def to_cuda(self): self.check_input() - if self.b_logits_rows is not None: - self.b_logits_rows = self.b_logits_rows.cuda(non_blocking=True) + if self.b_logits_row_idx is not None: + self.b_logits_row_idx = self.b_logits_row_idx.cuda(non_blocking=True) # Prefill 和 decode 都必须提供的公共张量。 self.b_req_idx = self.b_req_idx.cuda(non_blocking=True) @@ -213,7 +213,7 @@ class ModelOutput: # Immutable SSM configuration for this forward/graph; preserved when unpadding. ssm_run_config: Optional["SSMRunConfig"] = None # Number of physical decode rows represented by each selected logits row. - logits_row_stride: int = 1 + input_rows_per_logit: int = 1 def __post_init__(self) -> None: if self.mtp_collector is None: diff --git a/lightllm/common/basemodel/cuda_graph.py b/lightllm/common/basemodel/cuda_graph.py index 466dbfb952..5a38023b89 100644 --- a/lightllm/common/basemodel/cuda_graph.py +++ b/lightllm/common/basemodel/cuda_graph.py @@ -252,7 +252,7 @@ def replay(self, infer_state, infer_state1=None): return self._replay(infer_state) @torch.no_grad() - def warmup(self, model, logits_row_stride=1): + def warmup(self, model, input_rows_per_logit=1): logger.info("Begin capture cudagraph, use the --disable_cudagraph to disable it.") # for typing easy from .basemodel import TpPartBaseModel @@ -288,9 +288,9 @@ def warmup(self, model, logits_row_stride=1): multimodal_params=[{"images": [], "audios": []} for _ in range(batch_size)], **model._gen_special_model_input(batch_size), ) - if logits_row_stride > 1: - model_input.b_logits_rows = torch.arange( - logits_row_stride - 1, batch_size, logits_row_stride, dtype=torch.int64, device="cuda" + if input_rows_per_logit > 1: + model_input.b_logits_row_idx = torch.arange( + input_rows_per_logit - 1, batch_size, input_rows_per_logit, dtype=torch.int64, device="cuda" ) model_output: ModelOutput = model.forward(model_input) del model_output diff --git a/lightllm/common/basemodel/infer_struct.py b/lightllm/common/basemodel/infer_struct.py index 348a71f116..14f7c3f0d9 100755 --- a/lightllm/common/basemodel/infer_struct.py +++ b/lightllm/common/basemodel/infer_struct.py @@ -16,7 +16,7 @@ from .attention import BasePrefillAttState, BaseDecodeAttState if TYPE_CHECKING: - from lightllm.common.recurrent_state_types import SSMRunConfig + from lightllm.common.linear_state import SSMRunConfig class InferStateInfo: @@ -72,8 +72,8 @@ def __init__(self): # 在一些细节场景下需要有该信息区分一些资源的申请和管理。 self.microbatch_index: int = 0 self.ssm_run_config: Optional["SSMRunConfig"] = None - self.b_logits_rows: Optional[torch.Tensor] = None - self.logits_row_stride: int = 1 + self.b_logits_row_idx: Optional[torch.Tensor] = None + self.input_rows_per_logit: int = 1 # 当前 forward 独占的 hidden state 收集器。普通推理使用短生命周期实例； # Prefill CUDA Graph 使用随 graph infer state 长期保存的实例。 diff --git a/lightllm/common/basemodel/triton_kernel/linear_att/replayssm_compact.py b/lightllm/common/basemodel/triton_kernel/linear_att/compact_ssm.py similarity index 85% rename from lightllm/common/basemodel/triton_kernel/linear_att/replayssm_compact.py rename to lightllm/common/basemodel/triton_kernel/linear_att/compact_ssm.py index 113bb4e152..7173a3eed6 100644 --- a/lightllm/common/basemodel/triton_kernel/linear_att/replayssm_compact.py +++ b/lightllm/common/basemodel/triton_kernel/linear_att/compact_ssm.py @@ -9,7 +9,9 @@ import triton import triton.language as tl -from .ssm_autotune import configure_cache +from lightllm.common.linear_state import SSMRunConfig + +from .ssm_autotune import resolve_run_config @triton.jit @@ -115,13 +117,21 @@ def _compact( tl.store(sp, state, (kk[:, None] < K) & (vv[None, :] < V)) -class CompactSSMCache: +class CompactSSMBackend: def __init__( - self, state, verify_width, activation_dtype, *, num_key_heads=None, kda=False, lower_bound=-5.0, run_config=None + self, + state, + verify_width, + activation_dtype, + *, + num_key_heads=None, + kda=False, + gate_lower_bound=-5.0, + run_config=None, ): self.state = state self.kda = kda - self.lower_bound = lower_bound + self.gate_lower_bound = gate_lower_bound self.verify_width = verify_width # Verify and accepted-state reconstruction must share the same layout. self._config_is_fixed = run_config is not None @@ -132,7 +142,7 @@ def __init__( layers, slots, hv, k, v = state.shape self.num_key_heads = hv if num_key_heads is None else num_key_heads assert self.num_key_heads > 0 and hv % self.num_key_heads == 0 - self.hold = slots - 1 + self.hold_req_id = slots - 1 shape = (layers, slots, hv, verify_width) self.keys = torch.empty( (layers, slots, self.num_key_heads, verify_width, k), device=state.device, dtype=activation_dtype @@ -153,20 +163,23 @@ def merge_accepted_updates(self, reqs): # accept_updates already merged the accepted prefix into SSM state. pass - def snapshot_accepted_state(self, req_idx): + def get_accepted_state(self, req_idx): return self.state[:, req_idx] - def get_run_config(self, q, k, v, a, b, a_log, bias, cu_seqlens=None): - return configure_cache(self, "gdn", q, k, v, a, b, a_log, bias, cu_seqlens) + def get_run_config(self, q, k, v, a, b, a_log, bias, cu_seqlens=None) -> SSMRunConfig: + config = resolve_run_config(self, "compact", q, k, v, a, b, a_log, bias, cu_seqlens) + return SSMRunConfig(**config) def forward(self, linear_layer_idx, q, k, v, a, b, a_log, bias, reqs, cu_seqlens=None, run_config=None): assert cu_seqlens is not None, "compact replay is only used for speculative verify" config = run_config or self.get_run_config(q, k, v, a, b, a_log, bias, cu_seqlens) + if not isinstance(config, SSMRunConfig): + config = SSMRunConfig(**config) _, slots, hv, kd, vd = self.state.shape assert self.num_key_heads in (q.shape[-2], hv) out = torch.empty_like(v) # Match native reduction layouts: changing BV can alter BF16 rounding. - bv = config["BV"] + bv = config.BV _compact[(triton.cdiv(vd, bv), reqs.numel(), hv)]( q, k, @@ -196,23 +209,25 @@ def forward(self, linear_layer_idx, q, k, v, a, b, a_log, bias, reqs, cu_seqlens kd, vd, self.verify_width, - self.hold, + self.hold_req_id, triton.next_power_of_2(kd), bv, False, self.kda, - self.lower_bound, - num_warps=config["num_warps"], - num_stages=config.get("num_stages", 3), + self.gate_lower_bound, + num_warps=config.num_warps, + num_stages=config.num_stages, ) - return out + return out, config - def accept_updates(self, reqs, accepted, run_config=None): - """Merge this round's prefix; accepted contains per-request last accepted indexes.""" + def accept_updates(self, reqs, accepted_state_indices, run_config=None): + """Merge this round's prefix; accepted_state_indices contains per-request last accepted indexes.""" assert run_config is not None or not self.batch_configs, "compact acceptance requires the verify configuration" layers, slots, hv, kd, vd = self.state.shape config = run_config or self.run_config - bv = config["BV"] + if not isinstance(config, SSMRunConfig): + config = SSMRunConfig(**config) + bv = config.BV _compact[(triton.cdiv(vd, bv), reqs.numel(), layers * hv)]( None, None, @@ -228,7 +243,7 @@ def accept_updates(self, reqs, accepted, run_config=None): self.betas, reqs, None, - accepted, + accepted_state_indices, None, 0, 0, @@ -242,12 +257,12 @@ def accept_updates(self, reqs, accepted, run_config=None): kd, vd, self.verify_width, - self.hold, + self.hold_req_id, triton.next_power_of_2(kd), bv, True, self.kda, - self.lower_bound, - num_warps=config["num_warps"], - num_stages=config.get("num_stages", 3), + self.gate_lower_bound, + num_warps=config.num_warps, + num_stages=config.num_stages, ) diff --git a/lightllm/common/basemodel/triton_kernel/linear_att/mtp_state_params.py b/lightllm/common/basemodel/triton_kernel/linear_att/mtp_state_params.py index 685856fd0c..b0f2ea6b79 100644 --- a/lightllm/common/basemodel/triton_kernel/linear_att/mtp_state_params.py +++ b/lightllm/common/basemodel/triton_kernel/linear_att/mtp_state_params.py @@ -7,10 +7,8 @@ def _build_dynamic_mtp_linear_att_state_params_kernel( b_req_idx, b_mtp_index, - req_to_mtp_state_index, out_cu_q_seq_len, out_conv_buffer_idx, - out_num_accepted_tokens, batch_size, hold_req_id, SEQUENCE_CAPACITY: tl.constexpr, @@ -31,29 +29,20 @@ def _build_dynamic_mtp_linear_att_state_params_kernel( # Both the request-pool capacity and graph token bucket are static. tl.store(out_cu_q_seq_len + offsets, actual_token_num, mask=cu_mask) tl.store(out_conv_buffer_idx + offsets, hold_req_id, mask=sequence_mask) - tl.store(out_num_accepted_tokens + offsets, 1, mask=sequence_mask) tl.debug_barrier() # mtp_index restarts at zero on the first row of every request group. is_start = valid_row & (mtp_index == 0) sequence_index = tl.cumsum(tl.where(is_start, 1, 0), axis=0) - 1 - accepted_state_index = tl.load( - req_to_mtp_state_index + req_idx, - mask=is_start, - other=0, - ) - tl.store(out_cu_q_seq_len + sequence_index, offsets, mask=is_start) tl.store(out_conv_buffer_idx + sequence_index, req_idx, mask=is_start) - tl.store(out_num_accepted_tokens + sequence_index, accepted_state_index + 1, mask=is_start) def build_dynamic_mtp_linear_att_state_params( b_req_idx: torch.Tensor, b_mtp_index: torch.Tensor, - req_to_mtp_state_index: torch.Tensor, hold_req_id: int, -) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: +) -> tuple[torch.Tensor, torch.Tensor]: """Convert compact MTP verify rows to variable-length GDN sequences. Rows remain request-major after dynamic verification trimming, but each request can @@ -71,13 +60,11 @@ def build_dynamic_mtp_linear_att_state_params( request 4 -> rows [3, 4), query length 1 request 9 -> rows [4, 6), query length 2 - Runtime ``H`` rows are graph padding and do not form query sequences. If - ``req_to_mtp_state_index`` stores ``{7: 2, 4: 0, 9: 1}``, the fixed-shape - outputs are:: + Runtime ``H`` rows are graph padding and do not form query sequences. + The fixed-shape outputs are:: b1_cu_q_seq_len = [0, 3, 4, 6, 6, 6, 6, 6, 6] b_conv_buffer_idx = [7, 4, 9, H, H, H, H, H] - b_num_accepted_tokens = [3, 1, 2, 1, 1, 1, 1, 1] ``b_conv_buffer_idx`` therefore changes from one request id per query row to one request id per GDN sequence. Repeated cumulative lengths describe @@ -87,25 +74,22 @@ def build_dynamic_mtp_linear_att_state_params( one HOLD slot; the number of real sequences cannot exceed the other slots. """ - assert b_req_idx.is_cuda and b_mtp_index.is_cuda and req_to_mtp_state_index.is_cuda + assert b_req_idx.is_cuda and b_mtp_index.is_cuda assert b_req_idx.ndim == 1 and b_req_idx.shape == b_mtp_index.shape assert b_req_idx.dtype == torch.int32 and b_mtp_index.dtype == torch.int32 batch_size = b_req_idx.shape[0] assert batch_size > 0 - sequence_capacity = min(batch_size, req_to_mtp_state_index.numel() - 1) + sequence_capacity = min(batch_size, hold_req_id) assert sequence_capacity > 0 b1_cu_q_seq_len = torch.empty((sequence_capacity + 1,), dtype=torch.int32, device=b_req_idx.device) b_conv_buffer_idx = torch.empty((sequence_capacity,), dtype=torch.int32, device=b_req_idx.device) - b_num_accepted_tokens = torch.empty_like(b_conv_buffer_idx) _build_dynamic_mtp_linear_att_state_params_kernel[(1,)]( b_req_idx=b_req_idx, b_mtp_index=b_mtp_index, - req_to_mtp_state_index=req_to_mtp_state_index, out_cu_q_seq_len=b1_cu_q_seq_len, out_conv_buffer_idx=b_conv_buffer_idx, - out_num_accepted_tokens=b_num_accepted_tokens, batch_size=batch_size, hold_req_id=int(hold_req_id), SEQUENCE_CAPACITY=sequence_capacity, @@ -113,4 +97,4 @@ def build_dynamic_mtp_linear_att_state_params( num_warps=8, num_stages=1, ) - return b1_cu_q_seq_len, b_conv_buffer_idx, b_num_accepted_tokens + return b1_cu_q_seq_len, b_conv_buffer_idx diff --git a/lightllm/common/basemodel/triton_kernel/linear_att/replayssm.py b/lightllm/common/basemodel/triton_kernel/linear_att/replayssm.py index ec2e9f3d5e..fbedde9202 100644 --- a/lightllm/common/basemodel/triton_kernel/linear_att/replayssm.py +++ b/lightllm/common/basemodel/triton_kernel/linear_att/replayssm.py @@ -13,14 +13,16 @@ import triton import triton.language as tl -from .ssm_autotune import configure_cache +from lightllm.common.linear_state import SSMRunConfig + +from .ssm_autotune import resolve_run_config @triton.jit def _accept_updates( cursors, reqs, - accepted, + accepted_state_indices, N: tl.constexpr, HOLD: tl.constexpr, WIDTH: tl.constexpr, @@ -30,7 +32,7 @@ def _accept_updates( r = tl.load(reqs + i, i < N, HOLD) valid = (i < N) & (r != HOLD) cursor = tl.load(cursors + r, valid, 0) - count = tl.load(accepted + r, valid, 0) + 1 if WIDTH > 1 else 1 + count = tl.load(accepted_state_indices + r, valid, 0) + 1 if WIDTH > 1 else 1 tl.store(cursors + r, cursor + count, valid) @@ -218,7 +220,7 @@ def _replay( BK: tl.constexpr, BV: tl.constexpr, VARLEN: tl.constexpr, - PRECOMPUTE_STATE: tl.constexpr, + PRECOMPUTE_PROJECTIONS: tl.constexpr, BW: tl.constexpr, FUSED_FOLD: tl.constexpr, KDA: tl.constexpr, @@ -267,7 +269,7 @@ def _replay( tl.store(sp, state, (kk[:, None] < K) & (vv[None, :] < V)) base, n = L - base, 0 - if PRECOMPUTE_STATE: + if PRECOMPUTE_PROJECTIONS: projection_dtype: tl.constexpr = ( tl.bfloat16 if State.dtype.element_ty == tl.bfloat16 and Q.dtype.element_ty == tl.bfloat16 else tl.float32 ) @@ -326,7 +328,7 @@ def _replay( qk.to(projection_dtype), state.to(projection_dtype), input_precision=projection_precision ) - if PRECOMPUTE_STATE and not KDA: + if PRECOMPUTE_PROJECTIONS and not KDA: # Project the whole verify window against both accepted and new keys. # Only the small corrected-value solve below depends on token order. jj = tl.arange(0, BR) @@ -442,7 +444,7 @@ def _replay( tl.store(Betas + slot * (2 * L) + records, betas, fresh) return - if not PRECOMPUTE_STATE or KDA: + if not PRECOMPUTE_PROJECTIONS or KDA: history = ll < n keys = tl.load( Keys + (key_slot * (2 * L) + base) * K + ll[:, None] * K + kk[None, :], @@ -490,7 +492,7 @@ def _replay( keys *= tl.exp(g)[None, :] hk = tl.sum(keys * k[None, :], axis=1) hq = tl.sum(keys * q[None, :], axis=1) - if PRECOMPUTE_STATE: + if PRECOMPUTE_PROJECTIONS: checkpoint_sk = tl.sum(tl.where((tt == BW + t - start)[:, None], projections, 0), 0) checkpoint_sq = tl.sum(tl.where((tt == t - start)[:, None], projections, 0), 0) else: @@ -585,7 +587,7 @@ def _materialize( tl.store(out, state, (kk[:, None] < K) & (vv[None, :] < V)) -class ReplaySSMCache: +class ReplaySSMBackend: """Request-owned scratch; canonical CPU/PD checkpoints never include it.""" def __init__( @@ -597,17 +599,17 @@ def __init__( *, num_key_heads=None, kda=False, - lower_bound=-5.0, - projection_mode="inline", + gate_lower_bound=-5.0, + compute_mode="inline", run_config=None, fold_method="raw", ): assert state.dtype in (torch.float32, torch.bfloat16), "ReplaySSM requires FP32 or BF16 SSM state" assert verify_width > 0 and capacity >= max(4, verify_width) and capacity & (capacity - 1) == 0 - assert projection_mode in ("inline", "precompute", "parallel") + assert compute_mode in ("inline", "precompute", "parallel") assert fold_method in ("raw", "derived") - assert fold_method == "raw" or projection_mode == "parallel" - if projection_mode == "parallel": + assert fold_method == "raw" or compute_mode == "parallel" + if compute_mode == "parallel": assert not kda and verify_width > 1, "Parallel ReplaySSM requires GDN MTP" assert state.shape[-2:] == (128, 128), "Parallel ReplaySSM requires 128x128 heads" assert activation_dtype == torch.bfloat16, "Parallel ReplaySSM requires BF16 activations" @@ -615,13 +617,13 @@ def __init__( self.fold_method = fold_method self.state = state self.kda = kda - self.lower_bound = lower_bound + self.gate_lower_bound = gate_lower_bound self.fold_programs = ( 8 * torch.cuda.get_device_properties(state.device).multi_processor_count if state.is_cuda else 0 ) self.capacity = capacity self.verify_width = verify_width - self.projection_mode = projection_mode + self.compute_mode = compute_mode self.run_config = run_config self._config_is_fixed = run_config is not None self.batch_configs = {} @@ -629,7 +631,7 @@ def __init__( # Omission preserves the redundant layout for existing callers and A/B. self.num_key_heads = hv if num_key_heads is None else num_key_heads assert self.num_key_heads > 0 and hv % self.num_key_heads == 0 - self.hold = slots - 1 + self.hold_req_id = slots - 1 self.cursors = torch.zeros(slots, dtype=torch.int32, device=state.device) self.forward_cursors = torch.zeros_like(self.cursors) # Alternate halves on fold so CTAs cannot overwrite history another @@ -641,7 +643,7 @@ def __init__( activation_dtype if activation_dtype == torch.bfloat16 and verify_width > 1 - and projection_mode == "precompute" + and compute_mode == "precompute" and capacity <= 8 else torch.float32 ) @@ -678,12 +680,12 @@ def prepare_decode(self, reqs, cu_seqlens=None): cu_seqlens, active, reqs.numel(), - self.hold, + self.hold_req_id, self.capacity, cu_seqlens is not None, triton.next_power_of_2(reqs.numel()), self.verify_width, - self.projection_mode == "parallel", + self.compute_mode == "parallel", num_warps=4, ) programs = min(self.fold_programs, triton.cdiv(vd, 32) * reqs.numel() * layers * hv) @@ -720,14 +722,14 @@ def prepare_decode(self, reqs, cu_seqlens=None): ) return positions - def accept_updates(self, reqs, accepted=None): - """Accept this round's prefix; accepted contains per-request last accepted indexes.""" + def accept_updates(self, reqs, accepted_state_indices=None): + """Accept this round's prefix; accepted_state_indices contains per-request last accepted indexes.""" _accept_updates[(triton.cdiv(reqs.numel(), 256),)]( self.cursors, reqs, - accepted, + accepted_state_indices, reqs.numel(), - self.hold, + self.hold_req_id, self.verify_width, 256, ) @@ -758,7 +760,7 @@ def _materialize_accepted_state(self, reqs, output, snapshot): k, v, self.capacity, - self.hold, + self.hold_req_id, triton.next_power_of_2(k), bv, snapshot, @@ -767,7 +769,7 @@ def _materialize_accepted_state(self, reqs, output, snapshot): num_stages=config.get("num_stages", 3), ) - def snapshot_accepted_state(self, req_idx): + def get_accepted_state(self, req_idx): """Return a canonical checkpoint without changing the active state or history.""" layers, _, hv, k, v = self.state.shape output = torch.empty((layers, hv, k, v), dtype=self.state.dtype, device=self.state.device) @@ -775,21 +777,25 @@ def snapshot_accepted_state(self, req_idx): self._materialize_accepted_state(reqs, output, snapshot=True) return output - def get_run_config(self, q, k, v, a, b, a_log, bias, cu_seqlens=None): - if self.projection_mode == "parallel": - return self.run_config or {"BV": 128, "num_warps": 2, "num_stages": 1} - return configure_cache(self, "replay", q, k, v, a, b, a_log, bias, cu_seqlens) + def get_run_config(self, q, k, v, a, b, a_log, bias, cu_seqlens=None) -> SSMRunConfig: + if self.compute_mode == "parallel": + config = self.run_config or {"BV": 128, "num_warps": 2, "num_stages": 1} + else: + config = resolve_run_config(self, "replay", q, k, v, a, b, a_log, bias, cu_seqlens) + return SSMRunConfig(**config) def forward(self, linear_layer_idx, q, k, v, a, b, a_log, bias, reqs, cu_seqlens=None, run_config=None): config = run_config or self.get_run_config(q, k, v, a, b, a_log, bias, cu_seqlens) + if not isinstance(config, SSMRunConfig): + config = SSMRunConfig(**config) hv, kd, vd = self.state.shape[-3:] assert self.num_key_heads in (q.shape[-2], hv) axis = 1 if cu_seqlens is not None else 0 out = torch.empty_like(v) - bv = config["BV"] - precompute_state = self.verify_width > 1 and bv >= 16 and self.projection_mode in ("precompute", "parallel") - if self.projection_mode == "parallel": - assert precompute_state, "Parallel ReplaySSM requires BV >= 16" + bv = config.BV + precompute_projections = self.verify_width > 1 and bv >= 16 and self.compute_mode in ("precompute", "parallel") + if self.compute_mode == "parallel": + assert precompute_projections, "Parallel ReplaySSM requires BV >= 16" # Resolve tile dimensions in Python: bare max() in Triton JIT code # lowers to a tensor operation, which cannot define static shapes. bw = triton.next_power_of_2(self.verify_width) @@ -823,19 +829,19 @@ def forward(self, linear_layer_idx, q, k, v, a, b, a_log, bias, reqs, cu_seqlens kd, vd, self.capacity, - self.hold, + self.hold_req_id, triton.next_power_of_2(kd), bv, cu_seqlens is not None, - precompute_state, + precompute_projections, bw, - reqs.numel() > 12 and self.projection_mode != "parallel", + reqs.numel() > 12 and self.compute_mode != "parallel", self.kda, - self.lower_bound, - self.projection_mode == "parallel", + self.gate_lower_bound, + self.compute_mode == "parallel", BQK=max(16, 2 * bw), BR=triton.next_power_of_2(max(16, self.capacity + bw)), - num_warps=config["num_warps"], - num_stages=config.get("num_stages", 3), + num_warps=config.num_warps, + num_stages=config.num_stages, ) - return out + return out, config diff --git a/lightllm/common/basemodel/triton_kernel/linear_att/ssm_autotune.py b/lightllm/common/basemodel/triton_kernel/linear_att/ssm_autotune.py index 9f8618cc7d..052f61fb8f 100644 --- a/lightllm/common/basemodel/triton_kernel/linear_att/ssm_autotune.py +++ b/lightllm/common/basemodel/triton_kernel/linear_att/ssm_autotune.py @@ -1,4 +1,4 @@ -"""Tune a complete SSM update cycle without touching the serving state pool.""" +"""Tune one layer's SSM update cycle using the shared autotuner.""" import ast @@ -19,60 +19,63 @@ def get_configs(): ] -def static_key(cache, mode, q, k, v, a, b, cu_seqlens): +def static_key(ssm_backend, mode, q, k, v, a, b, cu_seqlens): axis = 1 if cu_seqlens is not None else 0 - layers, _, hv, kd, vd = cache.state.shape + _, _, hv, kd, vd = ssm_backend.state.shape # Keep filenames short, but distinguish packed projections, TP layouts, # state precision, recurrent rule and history capacity. key = { "mode": mode, - "v": 5, - "shape": str((layers, q.shape[-2], hv, cache.num_key_heads, kd, vd, cache.verify_width)), - "dtype": str((q.dtype, cache.state.dtype)), + "v": 6, + "shape": str((q.shape[-2], hv, ssm_backend.num_key_heads, kd, vd, ssm_backend.verify_width)), + "dtype": str((q.dtype, ssm_backend.state.dtype)), "strides": str((q.stride(axis), k.stride(axis), v.stride(axis), a.stride(0), b.stride(0))), - "history": cache.capacity if mode == "replay" else 0, - "gate": int(cache.kda), + "history": ssm_backend.capacity if mode == "replay" else 0, + "gate": int(ssm_backend.kda), "varlen": cu_seqlens is not None, } - if cache.kda: - key["lower"] = cache.lower_bound + if ssm_backend.kda: + key["lower"] = ssm_backend.gate_lower_bound if mode == "replay": - key["history_dtype"] = str(cache.keys.dtype) - if cache.projection_mode == "parallel": + key["history_dtype"] = str(ssm_backend.keys.dtype) + if ssm_backend.compute_mode == "parallel": key["projection"] = "parallel" - key["fold"] = cache.fold_method + key["fold"] = ssm_backend.fold_method return key key["projection"] = ( - ("kda_mma" if cache.kda else "gdn_mma") - if cache.projection_mode == "precompute" and cache.verify_width > 1 + ("kda_mma" if ssm_backend.kda else "gdn_mma") + if ssm_backend.compute_mode == "precompute" and ssm_backend.verify_width > 1 else "inline" ) return key -def run_key(cache, q, cu_seqlens): +def run_key(ssm_backend, q, cu_seqlens): tokens = q.shape[1 if cu_seqlens is not None else 0] return tokens, cu_seqlens.numel() - 1 if cu_seqlens is not None else tokens -def scratch_size_bytes(cache, batch, sequences, inputs): - # All cache tensors are request-owned; retain every layer and the HOLD slot. - slots = cache.state.shape[1] - pool_bytes = sum(x.numel() * x.element_size() for x in vars(cache).values() if isinstance(x, torch.Tensor)) +def scratch_size_bytes(ssm_backend, batch, sequences, inputs): + # State/history have a leading layer axis; cursors are shared across layers. + layers, slots = ssm_backend.state.shape[:2] + pool_bytes = sum( + x.numel() * x.element_size() // (layers if x.ndim > 1 else 1) + for x in vars(ssm_backend).values() + if isinstance(x, torch.Tensor) + ) input_bytes = sum( (1 + sum((size - 1) * stride for size, stride in zip(x.shape, x.stride()))) * x.element_size() for x in inputs if x.numel() ) - snapshot_bytes = cache.state.numel() * cache.state.element_size() // slots - return pool_bytes // slots * (batch + 1) + input_bytes + snapshot_bytes + 32 * (batch + sequences + 1) + return pool_bytes // slots * (batch + 1) + input_bytes + 32 * (batch + sequences + 1) -def rebuild_inputs(cache, mode, q, k, v, a, b, a_log, bias, cu_seqlens=None, workload=None, run_config=None): - tokens, sequences = run_key(cache, q, cu_seqlens) - batch = min(sequences, triton.cdiv(tokens, cache.verify_width)) if cu_seqlens is not None else tokens +def rebuild_inputs(ssm_backend, mode, q, k, v, a, b, a_log, bias, cu_seqlens=None, workload=None, run_config=None): + tokens, sequences = run_key(ssm_backend, q, cu_seqlens) + batch = min(sequences, triton.cdiv(tokens, ssm_backend.verify_width)) if cu_seqlens is not None else tokens if q.is_cuda: - required = scratch_size_bytes(cache, batch, sequences, (q, k, v, a, b, a_log, bias)) + required = scratch_size_bytes(ssm_backend, batch, sequences, (q, k, v, a, b, a_log, bias)) free, _ = torch.cuda.mem_get_info(q.device) available = free + torch.cuda.memory_reserved(q.device) - torch.cuda.memory_allocated(q.device) # Leave half for benchmark intermediates, CUDA Graphs and allocator fragmentation. @@ -80,28 +83,28 @@ def rebuild_inputs(cache, mode, q, k, v, a, b, a_log, bias, cu_seqlens=None, wor raise torch.cuda.OutOfMemoryError( f"SSM tuning scratch needs {required} bytes, exceeding its {available // 2}-byte budget" ) - layers, _, hv, kd, vd = cache.state.shape - state = torch.empty((layers, batch + 1, hv, kd, vd), device=q.device, dtype=cache.state.dtype).normal_(0, 0.02) + _, _, hv, kd, vd = ssm_backend.state.shape + state = torch.empty((1, batch + 1, hv, kd, vd), device=q.device, dtype=ssm_backend.state.dtype).normal_(0, 0.02) if mode == "replay": - scratch = type(cache)( + scratch = type(ssm_backend)( state, - cache.capacity, - cache.verify_width, + ssm_backend.capacity, + ssm_backend.verify_width, q.dtype, - num_key_heads=cache.num_key_heads, - kda=cache.kda, - lower_bound=cache.lower_bound, - projection_mode=cache.projection_mode, - fold_method=cache.fold_method, + num_key_heads=ssm_backend.num_key_heads, + kda=ssm_backend.kda, + gate_lower_bound=ssm_backend.gate_lower_bound, + compute_mode=ssm_backend.compute_mode, + fold_method=ssm_backend.fold_method, ) else: - scratch = type(cache)( + scratch = type(ssm_backend)( state, - cache.verify_width, + ssm_backend.verify_width, q.dtype, - num_key_heads=cache.num_key_heads, - kda=cache.kda, - lower_bound=cache.lower_bound, + num_key_heads=ssm_backend.num_key_heads, + kda=ssm_backend.kda, + gate_lower_bound=ssm_backend.gate_lower_bound, ) # Benchmark callbacks must not recursively select a configuration. scratch._config_is_fixed = True @@ -115,14 +118,13 @@ def rebuild_inputs(cache, mode, q, k, v, a, b, a_log, bias, cu_seqlens=None, wor torch.full((sequences - batch,), batch, dtype=torch.int32, device=q.device), ) ) - cu = (torch.arange(sequences + 1, dtype=torch.int32, device=q.device) * cache.verify_width).clamp_max(tokens) + cu = (torch.arange(sequences + 1, dtype=torch.int32, device=q.device) * ssm_backend.verify_width).clamp_max(tokens) lengths = cu[1 : batch + 1] - cu[:batch] accepted = torch.cat((lengths - 1, lengths.new_zeros(1))) partial = accepted.clone() - partial[:-1] = torch.minimum(reqs[:batch] % cache.verify_width, lengths - 1) - snapshot = torch.empty((layers, hv, kd, vd), device=q.device, dtype=state.dtype) + partial[:-1] = torch.minimum(reqs[:batch] % ssm_backend.verify_width, lengths - 1) return (), dict( - cache=scratch, + ssm_backend=scratch, mode=mode, q=inputs[0], k=inputs[1], @@ -132,7 +134,7 @@ def rebuild_inputs(cache, mode, q, k, v, a, b, a_log, bias, cu_seqlens=None, wor a_log=inputs[5], bias=inputs[6], cu_seqlens=cu if cu_seqlens is not None else None, - workload=(reqs, accepted, partial, snapshot), + workload=(reqs, accepted, partial), ) @@ -146,44 +148,39 @@ def rebuild_inputs(cache, mode, q, k, v, a, b, a_log, bias, cu_seqlens=None, wor rebuild_input_func=rebuild_inputs, warmup_all_exist_config=False, ) -def select_config(cache, mode, q, k, v, a, b, a_log, bias, cu_seqlens=None, workload=None, run_config=None): +def select_config(ssm_backend, mode, q, k, v, a, b, a_log, bias, cu_seqlens=None, workload=None, run_config=None): """Return a layout; only rebuilt disposable inputs execute the timed cycle.""" if workload is None: return run_config - cache.run_config = run_config - reqs, accepted, partial, snapshot = workload + ssm_backend.run_config = run_config + reqs, accepted, partial = workload if mode == "replay": - cache.cursors.zero_() + ssm_backend.cursors.zero_() # Include an entire history period and the subsequent fold. Repeating # just the first round would benchmark an empty history forever. tokens = q.shape[1 if cu_seqlens is not None else 0] - rounds = cache.capacity // min(cache.verify_width, tokens) + 1 + rounds = ssm_backend.capacity // min(ssm_backend.verify_width, tokens) + 1 else: - rounds = 2 # Partial acceptance and full acceptance, both across all layers. + rounds = 2 # Partial acceptance followed by full acceptance. for step in range(rounds): - cache.prepare_decode(reqs, cu_seqlens) - for linear_layer_idx in range(cache.state.shape[0]): - cache.forward(linear_layer_idx, q, k, v, a, b, a_log, bias, reqs, cu_seqlens) - cache.accept_updates(reqs, partial if mode != "replay" and step == 0 else accepted) - if mode == "replay" and step == rounds - 2: - cache._materialize_accepted_state(reqs[:1], snapshot, snapshot=True) - if mode == "replay": - cache.merge_accepted_updates(reqs) + ssm_backend.prepare_decode(reqs, cu_seqlens) + ssm_backend.forward(0, q, k, v, a, b, a_log, bias, reqs, cu_seqlens) + ssm_backend.accept_updates(reqs, partial if mode != "replay" and step == 0 else accepted) return run_config -def configure_cache(cache, mode, q, k, v, a, b, a_log, bias, cu_seqlens): - default = cache.run_config or {"BV": 32, "num_warps": 1} - if cache._config_is_fixed: +def resolve_run_config(ssm_backend, mode, q, k, v, a, b, a_log, bias, cu_seqlens): + default = ssm_backend.run_config or {"BV": 32, "num_warps": 1} + if ssm_backend._config_is_fixed: return default - key = (tuple(static_key(cache, mode, q, k, v, a, b, cu_seqlens).items()), run_key(cache, q, cu_seqlens)) - if key in cache.batch_configs: - return cache.batch_configs[key] + key = (tuple(static_key(ssm_backend, mode, q, k, v, a, b, cu_seqlens).items()), run_key(ssm_backend, q, cu_seqlens)) + if key in ssm_backend.batch_configs: + return ssm_backend.batch_configs[key] if not Autotuner.is_kernel_autotune_warmup(AutotuneKernelType.DECODE_ATTENTION): return default - config = select_config(cache, mode, q, k, v, a, b, a_log, bias, cu_seqlens) + config = select_config(ssm_backend, mode, q, k, v, a, b, a_log, bias, cu_seqlens) config = dict(config or default) # Each captured graph keeps its own layout, even after other buckets warm up. - cache.batch_configs[key] = config + ssm_backend.batch_configs[key] = config logger.info(f"SSM {mode} update-cycle config: {config}, warmup batch: {key[1]}") return config diff --git a/lightllm/common/basemodel/triton_kernel/linear_att_copy.py b/lightllm/common/basemodel/triton_kernel/linear_att_copy.py index 2d6438fd6e..45568703b6 100644 --- a/lightllm/common/basemodel/triton_kernel/linear_att_copy.py +++ b/lightllm/common/basemodel/triton_kernel/linear_att_copy.py @@ -25,9 +25,9 @@ def _copy_linear_att_state_to_kv_buffer( cpu_kv_ssm_stride_l, cpu_kv_ssm_stride_d, mtp_step, - conv_offsets, + accepted_state_indices, conv_element_bytes: tl.constexpr, - HAS_CONV_OFFSETS: tl.constexpr, + HAS_ACCEPTED_STATE_INDICES: 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, @@ -56,8 +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: - accepted_offset = tl.load(conv_offsets + cur_req_idx) + if HAS_ACCEPTED_STATE_INDICES: + accepted_offset = tl.load(accepted_state_indices + cur_req_idx) gpu_conv_base += accepted_offset * conv_element_bytes if mtp_step > 0: cur_state_req_idx += accepted_offset @@ -94,7 +94,7 @@ 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, + accepted_state_indices: torch.Tensor = None, ): # gpu_conv_state 的后两维可能是不连续的。 assert len(b_req_idx) == big_page_buffer_ids.shape[0] @@ -154,9 +154,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, + accepted_state_indices=accepted_state_indices, conv_element_bytes=conv_element_bytes, - HAS_CONV_OFFSETS=conv_offsets is not None, + HAS_ACCEPTED_STATE_INDICES=accepted_state_indices 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/basemodel/triton_kernel/post_process/apply_penalty_gpu_cache.py b/lightllm/common/basemodel/triton_kernel/post_process/apply_penalty_gpu_cache.py index f0fd2a4c78..a46fd7436c 100644 --- a/lightllm/common/basemodel/triton_kernel/post_process/apply_penalty_gpu_cache.py +++ b/lightllm/common/basemodel/triton_kernel/post_process/apply_penalty_gpu_cache.py @@ -3,7 +3,6 @@ import triton.language as tl import torch.nn.functional as F import numpy as np - from lightllm.common.req_manager import ReqSamplingParamsManager @@ -90,29 +89,27 @@ def apply_penalty_gpu_cache( b_mask_eos_reqs: torch.Tensor, eos_ids: torch.Tensor, sampling_params_manager: ReqSamplingParamsManager, - skip_token_penalties: bool = False, ): assert Logits.is_contiguous() - if not skip_token_penalties: - BLOCK_P = 2048 - num_warps = 8 - vocab_size = sampling_params_manager.vocab_size - req_to_out_token_id_counter = sampling_params_manager.req_to_out_token_id_counter - _fwd_kernel_apply_penalty_cache[(Logits.shape[0], triton.cdiv(vocab_size, BLOCK_P))]( - Logits=Logits, - stride_logit_b=Logits.stride(0), - stride_logit_s=Logits.stride(1), - b_req_idx=b_req_idx, - req_to_presence_penalty=sampling_params_manager.req_to_presence_penalty, - req_to_frequency_penalty=sampling_params_manager.req_to_frequency_penalty, - req_to_repetition_penalty=sampling_params_manager.req_to_repetition_penalty, - req_to_out_token_id_counter=req_to_out_token_id_counter, - stride_counter_r=req_to_out_token_id_counter.stride(0), - stride_counter_s=req_to_out_token_id_counter.stride(1), - vocab_size=vocab_size, - BLOCK_P=BLOCK_P, - num_warps=num_warps, - ) + BLOCK_P = 2048 + num_warps = 8 + vocab_size = sampling_params_manager.vocab_size + req_to_out_token_id_counter = sampling_params_manager.req_to_out_token_id_counter + _fwd_kernel_apply_penalty_cache[(Logits.shape[0], triton.cdiv(vocab_size, BLOCK_P))]( + Logits=Logits, + stride_logit_b=Logits.stride(0), + stride_logit_s=Logits.stride(1), + b_req_idx=b_req_idx, + req_to_presence_penalty=sampling_params_manager.req_to_presence_penalty, + req_to_frequency_penalty=sampling_params_manager.req_to_frequency_penalty, + req_to_repetition_penalty=sampling_params_manager.req_to_repetition_penalty, + req_to_out_token_id_counter=req_to_out_token_id_counter, + stride_counter_r=req_to_out_token_id_counter.stride(0), + stride_counter_s=req_to_out_token_id_counter.stride(1), + vocab_size=vocab_size, + BLOCK_P=BLOCK_P, + num_warps=num_warps, + ) BLOCK = 128 grid = (triton.cdiv(Logits.shape[0], BLOCK),) diff --git a/lightllm/common/basemodel/triton_kernel/post_process/argmax.py b/lightllm/common/basemodel/triton_kernel/post_process/argmax.py deleted file mode 100644 index d759a7dd8b..0000000000 --- a/lightllm/common/basemodel/triton_kernel/post_process/argmax.py +++ /dev/null @@ -1,47 +0,0 @@ -"""Two-stage vocabulary argmax with PyTorch tie and NaN semantics.""" - -import torch -import triton -import triton.language as tl - - -@triton.jit -def _argmax_partials(X, M, INDICES, NAN, V: tl.constexpr, S: tl.constexpr, P: tl.constexpr, B: tl.constexpr): - r = tl.program_id(0) - p = tl.program_id(1) - i = p * B + tl.arange(0, B) - x = tl.load(X + r * S + i, i < V, -float("inf")).to(tl.float32) - isnan = x != x - n = tl.min(tl.where((i < V) & isnan, i, 2147483647), 0) - m = tl.max(tl.where(isnan, -float("inf"), x), 0) - idx = tl.min(tl.where((i < V) & ~isnan & (x == m), i, 2147483647), 0) - tl.store(M + r * P + p, m) - tl.store(INDICES + r * P + p, idx) - tl.store(NAN + r * P + p, n) - - -@triton.jit -def _argmax_reduce(M, INDICES, NAN, OUTPUT, P: tl.constexpr, B: tl.constexpr): - r = tl.program_id(0) - i = tl.arange(0, B) - m = tl.load(M + r * P + i, i < P, -float("inf")) - idx = tl.load(INDICES + r * P + i, i < P, 2147483647) - nan = tl.min(tl.load(NAN + r * P + i, i < P, 2147483647), 0) - maximum = tl.max(m, 0) - winner = tl.min(tl.where(m == maximum, idx, 2147483647), 0) - tl.store(OUTPUT + r, tl.where(nan < 2147483647, nan, winner).to(tl.int64)) - - -def argmax(logits): - """Select the first maximum (or first NaN) from contiguous FP32 vocabulary rows.""" - assert logits.dtype == torch.float32 and logits.stride(1) == 1 - rows, vocab = logits.shape - block = 4096 - parts = triton.cdiv(vocab, block) - maxima = torch.empty((rows, parts), device=logits.device, dtype=torch.float32) - indices = torch.empty((rows, parts), device=logits.device, dtype=torch.int32) - nans = torch.empty_like(indices) - output = torch.empty(rows, device=logits.device, dtype=torch.int64) - _argmax_partials[(rows, parts)](logits, maxima, indices, nans, vocab, logits.stride(0), parts, block, num_warps=4) - _argmax_reduce[(rows,)](maxima, indices, nans, output, parts, triton.next_power_of_2(parts), num_warps=4) - return output diff --git a/lightllm/common/basemodel/triton_kernel/post_process/greedy.py b/lightllm/common/basemodel/triton_kernel/post_process/greedy.py deleted file mode 100644 index 827b540163..0000000000 --- a/lightllm/common/basemodel/triton_kernel/post_process/greedy.py +++ /dev/null @@ -1,66 +0,0 @@ -"""Greedy token selection and log-probability without a full softmax buffer.""" - -import torch -import triton -import triton.language as tl - - -@triton.jit -def _partials(X, T, M, S, INDICES, NAN, V: tl.constexpr, STRIDE: tl.constexpr, P: tl.constexpr, B: tl.constexpr): - row = tl.program_id(0) - part = tl.program_id(1) - i = part * B + tl.arange(0, B) - valid = i < V - x = tl.load(X + row * STRIDE + i, valid, other=-float("inf")).to(tl.float32) - temperature = tl.load(T + row) - x = tl.div_rn(x, temperature) - tl.store(X + row * STRIDE + i, x, valid) - is_nan = x != x - nan = tl.min(tl.where(valid & is_nan, i, 2147483647), 0) - maximum = tl.max(tl.where(is_nan, -float("inf"), x), 0) - winner = tl.min(tl.where(valid & ~is_nan & (x == maximum), i, 2147483647), 0) - denom = tl.sum(tl.where(valid & (maximum != -float("inf")), tl.exp(x - maximum), 0.0), 0) - offset = row * P + part - tl.store(M + offset, maximum) - tl.store(S + offset, denom) - tl.store(INDICES + offset, winner) - tl.store(NAN + offset, nan) - - -@triton.jit -def _finish(M, S, INDICES, NAN, OUT, LOGP, P: tl.constexpr, B: tl.constexpr): - row = tl.program_id(0) - i = tl.arange(0, B) - valid = i < P - m = tl.load(M + row * P + i, valid, other=-float("inf")) - s = tl.load(S + row * P + i, valid, other=0.0) - ids = tl.load(INDICES + row * P + i, valid, other=2147483647) - nans = tl.load(NAN + row * P + i, valid, other=2147483647) - maximum = tl.max(m, 0) - winner = tl.min(tl.where(valid & (m == maximum), ids, 2147483647), 0) - first_nan = tl.min(nans, 0) - winner = tl.where(first_nan < 2147483647, first_nan, winner) - denom = tl.sum(tl.where(valid, s * tl.exp(m - maximum), 0.0), 0) - logp = -tl.log(denom) - logp = tl.where(first_nan < 2147483647, float("nan"), logp) - tl.store(OUT + row, winner.to(tl.int64)) - tl.store(LOGP + row, logp) - - -def greedy_sample(logits, temperatures): - """Scale logits in place, preserving torch.div rounding and first-index argmax ties.""" - assert logits.dtype == torch.float32 and logits.stride(1) == 1 - rows, vocab = logits.shape - block = 4096 - parts = triton.cdiv(vocab, block) - maximum = torch.empty((rows, parts), device=logits.device, dtype=torch.float32) - sums = torch.empty_like(maximum) - indices = torch.empty((rows, parts), device=logits.device, dtype=torch.int32) - nans = torch.empty_like(indices) - ids = torch.empty(rows, device=logits.device, dtype=torch.int64) - logp = torch.empty(rows, device=logits.device, dtype=torch.float32) - _partials[(rows, parts)]( - logits, temperatures, maximum, sums, indices, nans, vocab, logits.stride(0), parts, block, num_warps=4 - ) - _finish[(rows,)](maximum, sums, indices, nans, ids, logp, parts, triton.next_power_of_2(parts), num_warps=4) - return ids, logp 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 315f9dfcda..2989c9a5cd 100644 --- a/lightllm/common/kv_cache_mem_manager/qwen3next_mem_manager.py +++ b/lightllm/common/kv_cache_mem_manager/qwen3next_mem_manager.py @@ -73,9 +73,9 @@ def _free_linear_att_buffers(self): return def write_to_shm(self, req_manager): - # RecurrentState contains only IPC-safe state/operator configuration; + # LinearStateManager contains only IPC-safe state/operator configuration; # worker execution resources must remain outside this shared object. - self.recurrent_state = req_manager.recurrent_state + self.linear_state_manager = req_manager.linear_state_manager # super().write_to_shm() 会用 ForkingPickler 序列化本对象，torch 在 dump 时会把 # CPU tensor 的 storage 原地迁到共享内存，使本进程大页 state cache 原本 # pinned(cudaHostAlloc) 的内存退化为普通 shm mmap，之后 Triton kernel 携带该指针 @@ -223,7 +223,7 @@ def get_dp_mems(self, mem_managers: List["Qwen3NextMemManager"], dp_index: int, dp_mems = mem_managers[(dp_index * dp_world_size) : ((dp_index + 1) * dp_world_size)] assert len(dp_mems) == dp_world_size for mem in dp_mems: - assert hasattr(mem, "recurrent_state") + assert hasattr(mem, "linear_state_manager") assert mem.linear_config.linear_layer_num == self.linear_config.linear_layer_num assert mem.linear_config.conv_kernel_size == self.linear_config.conv_kernel_size assert mem.linear_config.head_linear_k_dim == self.linear_config.head_linear_k_dim @@ -272,7 +272,7 @@ def _write_one_rank( conv_page: torch.Tensor, ssm_page: torch.Tensor, ): - conv_state, ssm_state = mem.recurrent_state.snapshot(req_idx) + conv_state, ssm_state = mem.linear_state_manager.get_accepted_state(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) return @@ -450,7 +450,7 @@ def _read_one_rank( conv_page: torch.Tensor, ssm_page: torch.Tensor, ): - conv_state, ssm_state = mem.recurrent_state.restore_target(req_idx) + conv_state, ssm_state = mem.linear_state_manager.prepare_restore(req_idx) self._copy_page_to_conv_state(conv_page, conv_state, mem, tp_index) self._copy_page_to_ssm_state(ssm_page, ssm_state, mem, tp_index) return diff --git a/lightllm/common/recurrent_state.py b/lightllm/common/linear_state.py similarity index 57% rename from lightllm/common/recurrent_state.py rename to lightllm/common/linear_state.py index d357c1ed06..1351d6bbb2 100644 --- a/lightllm/common/recurrent_state.py +++ b/lightllm/common/linear_state.py @@ -5,24 +5,36 @@ Attention backends retain responsibility for recurrent computation. """ +from dataclasses import dataclass from typing import TYPE_CHECKING, NamedTuple, Optional, Sequence, Tuple, Union import torch from lightllm.common.state_cache_manager import LayerCache, LinearAttCacheConfig -from lightllm.common.recurrent_state_types import SSMRunConfig, SSMUpdateBackend if TYPE_CHECKING: - from lightllm.common.basemodel.triton_kernel.linear_att.replayssm import ReplaySSMCache as _ReplaySSMCache - from lightllm.common.basemodel.triton_kernel.linear_att.replayssm_compact import CompactSSMCache as _CompactSSMCache + from lightllm.common.basemodel.triton_kernel.linear_att.replayssm import ReplaySSMBackend + from lightllm.common.basemodel.triton_kernel.linear_att.compact_ssm import CompactSSMBackend -class RecurrentStateSnapshot(NamedTuple): +@dataclass(frozen=True, slots=True) +class SSMRunConfig: + BV: int + num_warps: int + num_stages: int = 3 + + def __post_init__(self): + for value in (self.BV, self.num_warps, self.num_stages): + if type(value) is not int: + raise TypeError("SSM kernel configuration must be an integer") + + +class LinearStateTensors(NamedTuple): conv: torch.Tensor ssm: torch.Tensor -class RecurrentState: +class LinearStateManager: """Own request state and its lifecycle, without dispatching attention. PD currently shares this object through CUDA IPC as part of the memory @@ -36,8 +48,8 @@ class RecurrentState: def __init__(self, max_request_num, linear_config: LinearAttCacheConfig, args, *, kda=False, device="cuda"): self.linear_config = linear_config self.mtp_step = args.mtp_step - self.slots_per_req = self.mtp_step + 1 if args.ssm_state_mode == "native" else 1 - self.accepted_index = ( + self.ssm_slots_per_req = self.mtp_step + 1 if args.ssm_state_mode == "native" else 1 + self.accepted_state_index = ( torch.zeros(max_request_num + 1, dtype=torch.int32, device=device) if self.mtp_step > 0 else None ) self.conv_state = LayerCache( @@ -48,93 +60,101 @@ def __init__(self, max_request_num, linear_config: LinearAttCacheConfig, args, * device=device, ) self.ssm_state = LayerCache( - size=(max_request_num + 1) * self.slots_per_req, + size=(max_request_num + 1) * self.ssm_slots_per_req, dtype=linear_config.ssm_state_dtype, shape=linear_config.get_ssm_state_shape(), layer_num=linear_config.linear_layer_num, device=device, ) - self.updates: Optional[Union["_ReplaySSMCache", "_CompactSSMCache"]] = None + self.ssm_backend: Optional[Union["ReplaySSMBackend", "CompactSSMBackend"]] = None if args.ssm_state_mode == "replay": - from lightllm.common.basemodel.triton_kernel.linear_att.replayssm import ReplaySSMCache + from lightllm.common.basemodel.triton_kernel.linear_att import replayssm - self.updates = ReplaySSMCache( + self.ssm_backend = replayssm.ReplaySSMBackend( self.ssm_state.buffer, args.ssm_replay_config.get("cache_len", 16), self.mtp_step + 1, linear_config.conv_state_dtype, num_key_heads=linear_config.num_linear_k_heads, kda=kda, - lower_bound=linear_config.gate_lower_bound, - projection_mode=args.ssm_replay_config.get("compute_mode", "inline"), + gate_lower_bound=linear_config.gate_lower_bound, + compute_mode=args.ssm_replay_config.get("compute_mode", "inline"), fold_method=args.ssm_replay_config.get("fold_method", "raw"), ) elif args.ssm_state_mode == "compact": - from lightllm.common.basemodel.triton_kernel.linear_att.replayssm_compact import CompactSSMCache + from lightllm.common.basemodel.triton_kernel.linear_att import compact_ssm - self.updates = CompactSSMCache( + self.ssm_backend = compact_ssm.CompactSSMBackend( self.ssm_state.buffer, self.mtp_step + 1, linear_config.conv_state_dtype, num_key_heads=linear_config.num_linear_k_heads, kda=kda, - lower_bound=linear_config.gate_lower_bound, + gate_lower_bound=linear_config.gate_lower_bound, ) - @property - def update_backend(self) -> Optional[SSMUpdateBackend]: - """Optional compact/replay execution backend for attention to call.""" - return self.updates - - def _clear_positions(self, req_idx): - if self.updates is not None: - self.updates.clear_history(req_idx) - if self.accepted_index is not None: - self.accepted_index[req_idx].zero_() + def _reset_tracking(self, req_idx): + if self.ssm_backend is not None: + self.ssm_backend.clear_history(req_idx) + if self.accepted_state_index is not None: + self.accepted_state_index[req_idx].zero_() def reset(self, req_idx): """Initialize a newly allocated/reused request, including speculative slots.""" - self._clear_positions(req_idx) + self._reset_tracking(req_idx) self.conv_state.buffer[:, req_idx].zero_() - start = req_idx * self.slots_per_req - self.ssm_state.buffer[:, start : start + self.slots_per_req].zero_() + start = req_idx * self.ssm_slots_per_req + self.ssm_state.buffer[:, start : start + self.ssm_slots_per_req].zero_() def prepare_prefill(self, reqs): - if self.updates is not None: - self.updates.merge_accepted_updates(reqs) - return reqs * self.slots_per_req + """Merge accepted history and return convolution and SSM buffer indexes.""" + if self.ssm_backend is not None: + self.ssm_backend.merge_accepted_updates(reqs) + return reqs, reqs * self.ssm_slots_per_req def prepare_decode(self, reqs, cu_seqlens=None, *, single_token=False): - if self.updates is not None: - self.updates.prepare_decode(reqs, cu_seqlens) + if self.ssm_backend is not None: + self.ssm_backend.prepare_decode(reqs, cu_seqlens) if single_token: # Preserve the forward cursor snapshot before accepting the # guaranteed token, including when forward is graph-replayed. - self.updates.accept_updates(reqs) + self.ssm_backend.accept_updates(reqs) + + def prepare_mtp_decode(self, reqs, cu_seqlens): + """Return Conv/SSM indexes and prior accepted token counts; prepare replay history.""" + accepted_token_counts = self.accepted_state_index[reqs] + 1 + accepted_token_counts = torch.where(reqs == self.accepted_state_index.numel() - 1, 1, accepted_token_counts) + if self.ssm_backend is not None: + ssm_indexes = reqs + else: + offsets = torch.arange(self.ssm_slots_per_req, device=reqs.device, dtype=reqs.dtype) + ssm_indexes = reqs[:, None] * self.ssm_slots_per_req + offsets[None, :] + self.prepare_decode(reqs, cu_seqlens) + return reqs, ssm_indexes, accepted_token_counts - def commit( + def commit_accepted_state( self, b_req_mtp_start_loc, b_req_idx, b_mtp_index, accepted_index, verify_width, - batches: Optional[Sequence[Tuple[int, Optional[SSMRunConfig]]]] = None, + linear_state_batches: Optional[Sequence[Tuple[int, Optional[SSMRunConfig]]]] = None, ): """Commit accepted prefixes without requiring full-state materialization. - ``batches`` partitions logical requests in the same order as the MTP + ``linear_state_batches`` partitions logical requests in the same order as the MTP start locations. Each configuration belongs to the forward that wrote that partition's scratch; padding rows do not participate in acceptance. """ from lightllm.common.basemodel.triton_kernel.mtp_utils import linear_att_mtp_state_index_update - from lightllm.common.basemodel.triton_kernel.linear_att.replayssm import ReplaySSMCache + from lightllm.common.basemodel.triton_kernel.linear_att import replayssm replay_cursors = None - if isinstance(self.updates, ReplaySSMCache) and self.updates.verify_width > 1: - replay_cursors = self.updates.cursors + if isinstance(self.ssm_backend, replayssm.ReplaySSMBackend) and self.ssm_backend.verify_width > 1: + replay_cursors = self.ssm_backend.cursors linear_att_mtp_state_index_update( - req_to_mtp_state_index=self.accepted_index, + req_to_mtp_state_index=self.accepted_state_index, b_req_mtp_start_loc=b_req_mtp_start_loc, b_req_idx=b_req_idx, b_mtp_index=b_mtp_index, @@ -142,25 +162,25 @@ def commit( verify_width=verify_width, replay_cursors=replay_cursors, ) - if self.updates is None or replay_cursors is not None: + if self.ssm_backend is None or replay_cursors is not None: return reqs = b_req_idx[b_req_mtp_start_loc.long()] - if batches is None or isinstance(self.updates, ReplaySSMCache): - self.updates.accept_updates(reqs, self.accepted_index) + if linear_state_batches is None or isinstance(self.ssm_backend, replayssm.ReplaySSMBackend): + self.ssm_backend.accept_updates(reqs, self.accepted_state_index) return - assert sum(size for size, _ in batches) == reqs.numel() + assert sum(size for size, _ in linear_state_batches) == reqs.numel() start = 0 - for size, run_config in batches: + for size, run_config in linear_state_batches: if size: - self.updates.accept_updates( + self.ssm_backend.accept_updates( reqs[start : start + size], - self.accepted_index, + self.accepted_state_index, run_config=run_config, ) start += size - def snapshot(self, req_idx): - """Export the current accepted prefix without mutating active state. + def get_accepted_state(self, req_idx): + """Return the current accepted prefix without mutating active state. The caller associates this prefix with its cache/transport token boundary. These tensors may alias the active pool: enqueue consuming copies on the @@ -170,40 +190,40 @@ def snapshot(self, req_idx): """ width = self.linear_config.get_conv_state_shape()[-1] conv = self.conv_state.buffer[:, req_idx, ..., :width] - if self.accepted_index is not None: - offsets = torch.arange(width, device=conv.device) + self.accepted_index[req_idx] + if self.accepted_state_index is not None: + offsets = torch.arange(width, device=conv.device) + self.accepted_state_index[req_idx] conv = self.conv_state.buffer[:, req_idx].index_select(-1, offsets) - if self.updates is not None: - ssm = self.updates.snapshot_accepted_state(req_idx) - elif self.slots_per_req > 1: - index = req_idx * self.slots_per_req + self.accepted_index[req_idx] + if self.ssm_backend is not None: + ssm = self.ssm_backend.get_accepted_state(req_idx) + elif self.ssm_slots_per_req > 1: + index = req_idx * self.ssm_slots_per_req + self.accepted_state_index[req_idx] ssm = self.ssm_state.buffer.index_select(1, index.long().view(1))[:, 0] else: ssm = self.ssm_state.buffer[:, req_idx] - return RecurrentStateSnapshot(conv, ssm) + return LinearStateTensors(conv, ssm) - def restore_target(self, req_idx): + def prepare_restore(self, req_idx): """Invalidate old history and return canonical destinations for restore. The caller must fill both tensors before the next forward. This lets PD transport unpack TP shards directly into the pool without a staging copy. - The same stream/completion rules as ``snapshot`` apply. + The same stream/completion rules as ``get_accepted_state`` apply. """ - self._clear_positions(req_idx) + self._reset_tracking(req_idx) width = self.linear_config.get_conv_state_shape()[-1] - return RecurrentStateSnapshot( + return LinearStateTensors( self.conv_state.buffer[:, req_idx, ..., :width], - self.ssm_state.buffer[:, req_idx * self.slots_per_req], + self.ssm_state.buffer[:, req_idx * self.ssm_slots_per_req], ) def restore(self, req_idx, conv, ssm): - target = self.restore_target(req_idx) + target = self.prepare_restore(req_idx) target.conv.copy_(conv, non_blocking=True) target.ssm.copy_(ssm, non_blocking=True) def save_checkpoints(self, b_req_idx, req_indexes, buffer_indexes, destination): """Copy current accepted prefixes into caller-allocated checkpoint slots.""" - if self.updates is None: + if self.ssm_backend is None: from lightllm.common.basemodel.triton_kernel.linear_att_copy import copy_linear_att_state_to_kv_buffer indexes = torch.tensor(buffer_indexes, dtype=torch.int32, device="cpu").to( @@ -216,19 +236,19 @@ def save_checkpoints(self, b_req_idx, req_indexes, buffer_indexes, destination): gpu_ssm_state=self.ssm_state.buffer, cpu_kv_conv_state=destination.conv_state_cache.buffer, cpu_kv_ssm_state=destination.ssm_state_cache.buffer, - mtp_step=self.slots_per_req - 1, - conv_offsets=self.accepted_index, + mtp_step=self.ssm_slots_per_req - 1, + accepted_state_indices=self.accepted_state_index, ) else: for req_idx, buffer_idx in zip(req_indexes, buffer_indexes): if buffer_idx != -1: - self.export(req_idx, *destination.get_state_cache(buffer_idx)) + self.copy_accepted_state_to(req_idx, *destination.get_state_cache(buffer_idx)) - def export(self, req_idx, conv, ssm): - snapshot = self.snapshot(req_idx) - conv.copy_(snapshot.conv, non_blocking=True) - ssm.copy_(snapshot.ssm, non_blocking=True) + def copy_accepted_state_to(self, req_idx, conv, ssm): + accepted_state = self.get_accepted_state(req_idx) + conv.copy_(accepted_state.conv, non_blocking=True) + ssm.copy_(accepted_state.ssm, non_blocking=True) - def layer_buffers(self, layer_idx_in_all): + def get_layer_buffers(self, layer_idx_in_all): linear_layer_idx = self.linear_config.get_linear_layer_index(layer_idx_in_all) return self.conv_state.buffer[linear_layer_idx], self.ssm_state.buffer[linear_layer_idx] diff --git a/lightllm/common/recurrent_state_types.py b/lightllm/common/recurrent_state_types.py deleted file mode 100644 index ceb13e74f7..0000000000 --- a/lightllm/common/recurrent_state_types.py +++ /dev/null @@ -1,81 +0,0 @@ -"""Immutable execution metadata and the attention-facing SSM updater contract.""" - -from collections.abc import Iterator, Mapping -from dataclasses import dataclass -from typing import Optional, Protocol - -import torch - - -@dataclass(frozen=True, slots=True) -class SSMRunConfig(Mapping[str, int]): - """A detached kernel layout, also readable by existing Triton launch code.""" - - BV: int - num_warps: int - num_stages: int = 3 - - def __post_init__(self): - for key in self: - if type(self[key]) is not int: - raise TypeError(f"SSM kernel configuration {key} must be an integer") - - @classmethod - def from_mapping(cls, config: Mapping[str, int]) -> "SSMRunConfig": - unknown = config.keys() - {"BV", "num_warps", "num_stages"} - if unknown: - raise ValueError(f"Unknown SSM kernel configuration fields: {unknown}") - return cls(BV=config["BV"], num_warps=config["num_warps"], num_stages=config.get("num_stages", 3)) - - def __getitem__(self, key: str) -> int: - if key == "BV": - return self.BV - if key == "num_warps": - return self.num_warps - if key == "num_stages": - return self.num_stages - raise KeyError(key) - - def __iter__(self) -> Iterator[str]: - return iter(("BV", "num_warps", "num_stages")) - - def __len__(self) -> int: - return 3 - - -class SSMUpdateBackend(Protocol): - """Attention executes the updater directly; lifecycle stays with its owner. - - ``a`` and ``b`` are raw gate/beta projections. GDN uses scalar decay per - value head; KDA's ``a`` and ``bias`` include the key dimension. A missing - ``cu_seqlens`` denotes ordinary one-token-per-request decode. - """ - - def get_run_config( - self, - q: torch.Tensor, - k: torch.Tensor, - v: torch.Tensor, - a: torch.Tensor, - b: torch.Tensor, - a_log: torch.Tensor, - bias: torch.Tensor, - cu_seqlens: Optional[torch.Tensor] = None, - ) -> Mapping[str, int]: - ... - - def forward( - self, - linear_layer_idx: int, - q: torch.Tensor, - k: torch.Tensor, - v: torch.Tensor, - a: torch.Tensor, - b: torch.Tensor, - a_log: torch.Tensor, - bias: torch.Tensor, - reqs: torch.Tensor, - cu_seqlens: Optional[torch.Tensor] = None, - run_config: Optional[Mapping[str, int]] = None, - ) -> torch.Tensor: - ... diff --git a/lightllm/common/req_manager/hybrid_base.py b/lightllm/common/req_manager/hybrid_base.py index 370c8c4161..a1d178edec 100644 --- a/lightllm/common/req_manager/hybrid_base.py +++ b/lightllm/common/req_manager/hybrid_base.py @@ -66,7 +66,7 @@ def save_state(self, req_idx: int, buffer_idx: int, state_cache_manager): """请求 GPU 运行态 → 指定 CPU checkpoint 槽位；大小页共用，调用方负责分配槽位。""" def update_mtp_state( - self, b_req_mtp_start_loc, b_req_idx, b_mtp_index, accepted_index, verify_width, recurrent_state_batches=None + self, b_req_mtp_start_loc, b_req_idx, b_mtp_index, accepted_index, verify_width, linear_state_batches=None ): """接受推测 token 后更新运行态位置；需要调整状态索引的模型覆写。""" return diff --git a/lightllm/common/req_manager/linear_att.py b/lightllm/common/req_manager/linear_att.py index 9499ab2ee4..41f22dbfe4 100644 --- a/lightllm/common/req_manager/linear_att.py +++ b/lightllm/common/req_manager/linear_att.py @@ -2,7 +2,7 @@ import torch -from lightllm.common.recurrent_state import RecurrentState +from lightllm.common.linear_state import LinearStateManager from lightllm.common.state_cache_manager import LinearAttCacheConfig, LinearAttCacheManager from lightllm.utils.envs_utils import get_env_start_args @@ -26,30 +26,30 @@ def __init__(self, max_request_num, max_sequence_length, mem_manager, linear_con assert args.enable_prefill_decode_mixed is False self.big_page_token_num = args.linear_att_page_block_num * args.linear_att_hash_page_size self.linear_config = linear_config - self.recurrent_state = RecurrentState(max_request_num, linear_config, args, kda=self.kda) + self.linear_state_manager = LinearStateManager(max_request_num, linear_config, args, kda=self.kda) def init_hybrid_attention_state(self, req: "InferReq"): - self.recurrent_state.reset(req.req_idx) + self.linear_state_manager.reset(req.req_idx) def create_small_page_cache_manager(self, size: int): self.small_page_buffers = LinearAttCacheManager(size=size, linear_config=self.linear_config) return self.small_page_buffers def save_big_page_states(self, b_req_idx: torch.Tensor, req_indexes: List[int], buffer_indexes: List[int]): - self.recurrent_state.save_checkpoints(b_req_idx, req_indexes, buffer_indexes, self.big_page_buffers) + self.linear_state_manager.save_checkpoints(b_req_idx, req_indexes, buffer_indexes, self.big_page_buffers) def save_state(self, req_idx: int, buffer_idx: int, state_cache_manager: LinearAttCacheManager): - self.recurrent_state.export(req_idx, *state_cache_manager.get_state_cache(buffer_idx)) + self.linear_state_manager.copy_accepted_state_to(req_idx, *state_cache_manager.get_state_cache(buffer_idx)) def get_mamba_cache(self, layer_idx_in_all: int): - return self.recurrent_state.layer_buffers(layer_idx_in_all) + return self.linear_state_manager.get_layer_buffers(layer_idx_in_all) def update_mtp_state( - self, b_req_mtp_start_loc, b_req_idx, b_mtp_index, accepted_index, verify_width, recurrent_state_batches=None + self, b_req_mtp_start_loc, b_req_idx, b_mtp_index, accepted_index, verify_width, linear_state_batches=None ): - self.recurrent_state.commit( - b_req_mtp_start_loc, b_req_idx, b_mtp_index, accepted_index, verify_width, recurrent_state_batches + self.linear_state_manager.commit_accepted_state( + b_req_mtp_start_loc, b_req_idx, b_mtp_index, accepted_index, verify_width, linear_state_batches ) def restore_state(self, req: "InferReq", state_cache_manager: LinearAttCacheManager, buffer_idx: int): - self.recurrent_state.restore(req.req_idx, *state_cache_manager.get_state_cache(buffer_idx)) + self.linear_state_manager.restore(req.req_idx, *state_cache_manager.get_state_cache(buffer_idx)) diff --git a/lightllm/models/llama/layer_infer/post_layer_infer.py b/lightllm/models/llama/layer_infer/post_layer_infer.py index c186877eea..02cb01ebc9 100644 --- a/lightllm/models/llama/layer_infer/post_layer_infer.py +++ b/lightllm/models/llama/layer_infer/post_layer_infer.py @@ -243,8 +243,8 @@ def _draft_lm_head_and_gather( ) -> PostLayerOutput: """按 draft 配置返回候选 logits，未开启时回退到完整词表。""" - if infer_state.b_logits_rows is not None: - hidden = hidden.index_select(0, infer_state.b_logits_rows) + if infer_state.b_logits_row_idx is not None: + hidden = hidden.index_select(0, infer_state.b_logits_row_idx) token_num = hidden.shape[0] top_k = get_env_start_args().draft_vocab_topk_sampling if top_k is None: diff --git a/lightllm/server/api_start.py b/lightllm/server/api_start.py index 85cd8e7e6c..840537d794 100644 --- a/lightllm/server/api_start.py +++ b/lightllm/server/api_start.py @@ -44,7 +44,7 @@ def _set_envs_and_config(args: StartArgs): mp.set_start_method("spawn", force=True) -def _validate_ssm_state_mode(args: StartArgs): +def _validate_and_normalize_ssm_config(args: StartArgs): assert args.ssm_state_mode in ("native", "compact", "replay"), "Invalid ssm_state_mode" config = args.ssm_replay_config assert isinstance(config, dict), "ssm_replay_config must be a JSON object" @@ -363,7 +363,7 @@ def _launch_subprocesses(args: StartArgs): f"but got {args.batch_max_tokens}, {args.chunked_prefill_size}" ) - _validate_ssm_state_mode(args) + _validate_and_normalize_ssm_config(args) args.disable_chunked_mm_input = args.disable_chunked_mm_input or get_default_disable_chunked_mm_input( args.model_dir diff --git a/lightllm/server/router/model_infer/mode_backend/base_backend.py b/lightllm/server/router/model_infer/mode_backend/base_backend.py index 1c3803ce9b..7997f49934 100644 --- a/lightllm/server/router/model_infer/mode_backend/base_backend.py +++ b/lightllm/server/router/model_infer/mode_backend/base_backend.py @@ -1,8 +1,10 @@ import numpy as np import torch +import time +import threading import torch.distributed as dist +from typing import List, Tuple, Callable, Optional, Union from transformers.configuration_utils import PretrainedConfig - from lightllm.utils.infer_utils import set_random_seed from lightllm.utils.log_utils import init_logger from lightllm.models import get_draft_model_class, get_model @@ -760,6 +762,7 @@ def _get_classed_reqs( can_alloc_token_num = g_infer_context.get_can_alloc_token_num() for req_obj in ready_reqs: + if req_obj.filter_mark: finished_reqs.append(req_obj) continue @@ -959,17 +962,6 @@ def _trans_req_ids_to_req_objs(self, req_ids: List[int]) -> List[InferReq]: return [g_infer_context.requests_mapping[req_id] for req_id in req_ids] def _gen_argmax_token_ids(self, model_output: ModelOutput): - if ( - model_output.logits_token_ids is None - and model_output.logits.is_cuda - and model_output.logits.dtype == torch.float32 - and model_output.logits.shape[0] > 0 - and model_output.logits.stride(1) == 1 - and model_output.logits.shape[1] >= 4096 - ): - from lightllm.common.basemodel.triton_kernel.post_process.argmax import argmax - - return argmax(model_output.logits) token_indices = torch.argmax(model_output.logits, dim=-1) if model_output.logits_token_ids is not None: # 候选排列保证分数相同时优先选择最小的全局 token ID。 diff --git a/lightllm/server/router/model_infer/mode_backend/chunked_prefill/impl.py b/lightllm/server/router/model_infer/mode_backend/chunked_prefill/impl.py index 7cefa1b677..ac38e4e975 100644 --- a/lightllm/server/router/model_infer/mode_backend/chunked_prefill/impl.py +++ b/lightllm/server/router/model_infer/mode_backend/chunked_prefill/impl.py @@ -298,7 +298,7 @@ def decode_mtp( b_req_idx=model_input.b_req_idx, b_req_mtp_start_loc=b_req_mtp_start_loc, b_mtp_index=model_input.b_mtp_index, - recurrent_state_batches=((req_num, model_output.ssm_run_config),), + linear_state_batches=((req_num, model_output.ssm_run_config),), ) accepted_index_cpu = g_pin_mem_manager.async_copy_from_gpu_tensor( key="accepted_index", diff --git a/lightllm/server/router/model_infer/mode_backend/dp_backend/impl.py b/lightllm/server/router/model_infer/mode_backend/dp_backend/impl.py index f68464e719..ced5d1d64b 100644 --- a/lightllm/server/router/model_infer/mode_backend/dp_backend/impl.py +++ b/lightllm/server/router/model_infer/mode_backend/dp_backend/impl.py @@ -533,7 +533,7 @@ def decode_mtp(self, event_pack: OverlapEventPack, decode_reqs: List[InferReq]): b_req_idx=model_input.b_req_idx, b_req_mtp_start_loc=b_req_mtp_start_loc, b_mtp_index=model_input.b_mtp_index, - recurrent_state_batches=((req_num, model_output.ssm_run_config),), + linear_state_batches=((req_num, model_output.ssm_run_config),), ) accepted_index_cpu = g_pin_mem_manager.async_copy_from_gpu_tensor( key="accepted_index", @@ -808,7 +808,7 @@ def decode_overlap_mtp(self, event_pack: OverlapEventPack, decode_reqs: List[Inf b_req_idx=b_req_idx, b_req_mtp_start_loc=b_req_mtp_start_loc, b_mtp_index=b_mtp_index, - recurrent_state_batches=( + linear_state_batches=( (real_request_num0, model_output0.ssm_run_config), (real_request_num1, model_output1.ssm_run_config), ), diff --git a/lightllm/server/router/model_infer/mode_backend/generic_post_process.py b/lightllm/server/router/model_infer/mode_backend/generic_post_process.py index 68c6822645..90c96ecf2f 100644 --- a/lightllm/server/router/model_infer/mode_backend/generic_post_process.py +++ b/lightllm/server/router/model_infer/mode_backend/generic_post_process.py @@ -1,12 +1,9 @@ -from typing import List, Optional, Tuple - import torch - +from typing import List, Optional, Tuple from lightllm.common.basemodel.triton_kernel.post_process.apply_penalty import apply_penalty from lightllm.common.basemodel.triton_kernel.post_process.apply_penalty_gpu_cache import apply_penalty_gpu_cache from lightllm.common.basemodel.triton_kernel.post_process.apply_invalid_token import apply_invalid_token_ids from lightllm.common.basemodel.triton_kernel.post_process.apply_constraint_mask import apply_constraint_mask -from lightllm.common.basemodel.triton_kernel.post_process.greedy import greedy_sample from lightllm.server.router.model_infer.infer_batch import InferReq, g_infer_context from lightllm.server.router.model_infer.pin_mem_manager import g_pin_mem_manager from lightllm.utils.envs_utils import get_env_start_args @@ -35,7 +32,6 @@ def sample( skip_top_p, exist_req_use_random_seed, has_output_constraints, - skip_token_penalties, ) = _get_post_sample_tensors(reqs, sampling_params_manager, has_output) eos_ids = g_pin_mem_manager.gen_from_list(key="eos_ids", data=eos_id, dtype=torch.int32).cuda(non_blocking=True) @@ -77,7 +73,6 @@ def sample( b_mask_eos_reqs=b_mask_eos_reqs, eos_ids=eos_ids, sampling_params_manager=sampling_params_manager, - skip_token_penalties=skip_token_penalties, ) if has_invalid_token_ids: @@ -87,15 +82,6 @@ def sample( cu_invalid_token_num=cu_invalid_token_num, ) - if ( - is_all_greedy - and not has_output_constraints - and logits.is_cuda - and logits.dtype == torch.float32 - and logits.stride(1) == 1 - ): - return greedy_sample(logits, b_temperatures) - logits.div_(b_temperatures.view((-1, 1))) # Apply hard constraints after all penalties. In particular, the EOS length # penalty must not turn a masked -inf into NaN or re-enable a forbidden EOS. @@ -208,7 +194,6 @@ def _get_post_sample_tensors(reqs: List[InferReq], sampling_params_manager, has_ request_group=request_group, sampling_params_manager=sampling_params_manager, ) - skip_token_penalties = True # invalid token ids invalid_token_ids: List[int] = [] @@ -219,11 +204,6 @@ def _get_post_sample_tensors(reqs: List[InferReq], sampling_params_manager, has_ for i, req_obj in enumerate(reqs): sample_param = req_obj.sampling_param shm_param = sample_param.shm_param - skip_token_penalties = skip_token_penalties and ( - shm_param.presence_penalty == 0.0 - and shm_param.frequency_penalty == 0.0 - and shm_param.repetition_penalty == 1.0 - ) exponential_decay_length_penalty = shm_param.exponential_decay_length_penalty.to_tuple() out_token_len = req_obj.get_cur_total_len() - req_obj.shm_req.input_len length_penalty_param.append(max(out_token_len - exponential_decay_length_penalty[0], 0)) @@ -280,7 +260,6 @@ def _get_post_sample_tensors(reqs: List[InferReq], sampling_params_manager, has_ skip_top_p, exist_req_use_random_seed, has_output_constraints, - skip_token_penalties, ) diff --git a/lightllm/server/router/model_infer/mtp_speculative/proposers/eagle_with_att.py b/lightllm/server/router/model_infer/mtp_speculative/proposers/eagle_with_att.py index d5dcdca05b..4f84b7e8c7 100644 --- a/lightllm/server/router/model_infer/mtp_speculative/proposers/eagle_with_att.py +++ b/lightllm/server/router/model_infer/mtp_speculative/proposers/eagle_with_att.py @@ -73,8 +73,8 @@ def propose_next( verify_draft_input = copy.copy(target_model_input) verify_draft_input.input_ids = target_next_token_ids verify_draft_input.mtp_draft_input_hiddens = target_model_output.mtp_collector.spec_hidden - if draft_model.use_replayssm_draft_tail(): - verify_draft_input.b_logits_rows = accepted_tail_rows + if draft_model.should_project_draft_tail_logits(): + verify_draft_input.b_logits_row_idx = accepted_tail_rows extend_output = draft_model.forward(verify_draft_input) # 只在 req_num 行 logits 上进行 argmax，避免为未接受的 verify 行执行 @@ -82,13 +82,13 @@ def propose_next( accepted_tail_output = ModelOutput( logits=( extend_output.logits[:req_num] - if extend_output.logits_row_stride > 1 + if extend_output.input_rows_per_logit > 1 else extend_output.logits.index_select(0, accepted_tail_rows) ), logits_token_ids=( ( extend_output.logits_token_ids[:req_num] - if extend_output.logits_row_stride > 1 + if extend_output.input_rows_per_logit > 1 else extend_output.logits_token_ids.index_select(0, accepted_tail_rows) ) if extend_output.logits_token_ids is not None diff --git a/lightllm/server/router/model_infer/mtp_speculative/utils.py b/lightllm/server/router/model_infer/mtp_speculative/utils.py index be9a40c457..7b68d06179 100644 --- a/lightllm/server/router/model_infer/mtp_speculative/utils.py +++ b/lightllm/server/router/model_infer/mtp_speculative/utils.py @@ -22,7 +22,7 @@ def verify_mtp_tokens( b_req_idx: torch.Tensor, b_req_mtp_start_loc: torch.Tensor, b_mtp_index: torch.Tensor, - recurrent_state_batches=None, + linear_state_batches=None, ) -> Tuple[torch.Tensor, torch.Tensor]: """Verify target tokens and update recurrent MTP state when required.""" # 保留原调用接口中的步序参数；当前状态提交直接使用验证返回的接受长度。 @@ -40,7 +40,7 @@ def verify_mtp_tokens( b_mtp_index=b_mtp_index, accepted_index=accepted_index, verify_width=backend.max_draft_step + 1, - recurrent_state_batches=recurrent_state_batches, + linear_state_batches=linear_state_batches, ) return accept_lengths, accepted_index diff --git a/test/benchmark/kernels/benchmark_kda_ssm.py b/test/benchmark/kernels/benchmark_kda_ssm.py index 4a166228e2..8a8c4cfe58 100644 --- a/test/benchmark/kernels/benchmark_kda_ssm.py +++ b/test/benchmark/kernels/benchmark_kda_ssm.py @@ -14,8 +14,8 @@ import triton from lightllm.common.basemodel.triton_kernel.linear_att.fla.ops.kda_decode import fused_recurrent_kda -from lightllm.common.basemodel.triton_kernel.linear_att.replayssm import ReplaySSMCache -from lightllm.common.basemodel.triton_kernel.linear_att.replayssm_compact import CompactSSMCache +from lightllm.common.basemodel.triton_kernel.linear_att.replayssm import ReplaySSMBackend +from lightllm.common.basemodel.triton_kernel.linear_att.compact_ssm import CompactSSMBackend from lightllm.common.basemodel.triton_kernel.linear_att.ssm_autotune import get_configs @@ -35,19 +35,19 @@ def workload(args, batch, width, dtype, mode, acceptance): snapshot = torch.empty(layers, heads, dim, dim, dtype=dtype, device="cuda") slots = (batch + 1) * (width if mode == "native" else 1) state = torch.randn(layers, slots, heads, dim, dim, dtype=dtype, device="cuda") * 0.01 - cache = None + ssm_backend = None if mode == "compact": - cache = CompactSSMCache(state, width, q.dtype, kda=True) + ssm_backend = CompactSSMBackend(state, width, q.dtype, kda=True) elif mode != "native": - cache = ReplaySSMCache(state, 8, width, q.dtype, kda=True, projection_mode=mode) - if cache is not None: - cache._config_is_fixed = True + ssm_backend = ReplaySSMBackend(state, 8, width, q.dtype, kda=True, compute_mode=mode) + if ssm_backend is not None: + ssm_backend._config_is_fixed = True indexes = reqs[:, None] * width + torch.arange(width, device="cuda", dtype=torch.int32) if width == 1: indexes = reqs bytes_used = ( - sum(x.numel() * x.element_size() for x in vars(cache).values() if isinstance(x, torch.Tensor)) - if cache + sum(x.numel() * x.element_size() for x in vars(ssm_backend).values() if isinstance(x, torch.Tensor)) + if ssm_backend else state.numel() * state.element_size() ) @@ -55,10 +55,10 @@ def workload(args, batch, width, dtype, mode, acceptance): def cycle(): for _ in range(args.rounds): - if cache is not None: - cache.prepare_decode(reqs, cu) + if ssm_backend is not None: + ssm_backend.prepare_decode(reqs, cu) for layer in range(layers): - if cache is None: + if ssm_backend is None: fused_recurrent_kda( q, k, @@ -73,16 +73,16 @@ def cycle(): num_accepted_tokens=counts if cu is not None else None, ) else: - cache.forward(layer, q, k, v, gate, beta, alog, bias, reqs, cu) - if cache is not None: - cache.accept_updates(reqs, accepted) - if isinstance(cache, ReplaySSMCache): - cache._materialize_accepted_state(reqs[:1], snapshot, snapshot=True) - cache.merge_accepted_updates(reqs) + ssm_backend.forward(layer, q, k, v, gate, beta, alog, bias, reqs, cu) + if ssm_backend is not None: + ssm_backend.accept_updates(reqs, accepted) + if isinstance(ssm_backend, ReplaySSMBackend): + ssm_backend._materialize_accepted_state(reqs[:1], snapshot, snapshot=True) + ssm_backend.merge_accepted_updates(reqs) else: - snapshot.copy_(state[:, native_index] if cache is None else state[:, 0]) + snapshot.copy_(state[:, native_index] if ssm_backend is None else state[:, 0]) - return cache, cycle, bytes_used + return ssm_backend, cycle, bytes_used def main(): @@ -111,15 +111,17 @@ def main(): for width in args.width: for acceptance in ["full", "partial"] if width > 1 else ["full"]: for mode in ["native", "compact", "inline", "precompute"] if width > 1 else ["native", "inline"]: - cache, cycle, size = workload(args, batch, width, getattr(torch, dtype_name), mode, acceptance) + ssm_backend, cycle, size = workload( + args, batch, width, getattr(torch, dtype_name), mode, acceptance + ) candidates = [] - if cache is not None: + if ssm_backend is not None: for config in get_configs() if args.tune else [{"BV": 32, "num_warps": 2, "num_stages": 1}]: - cache.run_config = config + ssm_backend.run_config = config cycle() ms = triton.testing.do_bench_cudagraph(cycle, rep=50) candidates.append(dict(config=config, ms=ms)) - cache.run_config = min(candidates, key=lambda r: r["ms"])["config"] + ssm_backend.run_config = min(candidates, key=lambda r: r["ms"])["config"] cycle() samples = [ triton.testing.do_bench_cudagraph(cycle, rep=100) / args.rounds for _ in range(args.repeats) @@ -133,13 +135,13 @@ def main(): state_history_bytes=size, ms_per_round=statistics.median(samples), samples_ms=samples, - config=cache.run_config if cache else None, + config=ssm_backend.run_config if ssm_backend else None, candidates=candidates, ) rows.append(row) args.output.write_text(json.dumps(report, indent=2) + "\n") print(json.dumps(row), flush=True) - del cache, cycle + del ssm_backend, cycle gc.collect() torch.cuda.empty_cache() diff --git a/test/benchmark/kernels/benchmark_replayssm_service.py b/test/benchmark/kernels/benchmark_replayssm_service.py index 08a694a684..a1cee4826a 100644 --- a/test/benchmark/kernels/benchmark_replayssm_service.py +++ b/test/benchmark/kernels/benchmark_replayssm_service.py @@ -39,7 +39,7 @@ def prepare_prompt(index, rep): return prefix + suffix -def one(index, rep, ids, barrier, cohort_start): +def stream_request(index, rep, ids, barrier, cohort_start): payload = { "inputs": ids, "parameters": { @@ -83,12 +83,12 @@ def one(index, rep, ids, barrier, cohort_start): ) -def run(batch, rep): +def run_concurrent_requests(batch, rep): prompts = [prepare_prompt(i, rep) for i in range(batch)] barrier = threading.Barrier(batch) start = time.perf_counter() with concurrent.futures.ThreadPoolExecutor(max_workers=batch) as pool: - samples = list(pool.map(lambda i: one(i, rep, prompts[i], barrier, start), range(batch))) + samples = list(pool.map(lambda i: stream_request(i, rep, prompts[i], barrier, start), range(batch))) elapsed = time.perf_counter() - start begin = max(r["token_times_s"][0] for r in samples) + 0.5 end = min(r["token_times_s"][-1] for r in samples) - 0.5 @@ -109,11 +109,11 @@ def run(batch, rep): # Warm every measured graph bucket before timing, with the same requests on both revisions. for batch in args.batch: for warm in range(args.warmups): - run(batch, 99 + warm) + run_concurrent_requests(batch, 99 + warm) rows = [] for batch in args.batch: for rep in range(args.repeats): - row = run(batch, rep) + row = run_concurrent_requests(batch, rep) rows.append(row) Path(args.output).write_text(json.dumps(rows, indent=2)) print(json.dumps({k: v for k, v in row.items() if k != "requests"}), flush=True) diff --git a/unit_tests/common/basemodel/test_draft_tail_graph.py b/unit_tests/common/basemodel/test_draft_tail_graph.py index 75bdf67955..57ce202959 100644 --- a/unit_tests/common/basemodel/test_draft_tail_graph.py +++ b/unit_tests/common/basemodel/test_draft_tail_graph.py @@ -70,16 +70,16 @@ def replay(): def forward(state): hidden = state.input_ids[:, None].float() - logits = hidden if state.b_logits_rows is None else hidden.index_select(0, state.b_logits_rows) + logits = hidden if state.b_logits_row_idx is None else hidden.index_select(0, state.b_logits_row_idx) return ModelOutput( logits=logits, - logits_row_stride=state.logits_row_stride, + input_rows_per_logit=state.input_rows_per_logit, mtp_collector=ModelMtpOutputCollector(spec_hidden=hidden), ) model._token_forward = forward model._init_cudagraph() - assert warmups == [(model.draft_tail_graph, {"logits_row_stride": 5}), (model.graph, {})] + assert warmups == [(model.draft_tail_graph, {"input_rows_per_logit": 5}), (model.graph, {})] return model @@ -97,7 +97,7 @@ def make_input(rows, accepted=None): b_shared_seq_len=torch.zeros(rows, dtype=torch.int32), b_shared_radix_node_id=torch.zeros(rows, dtype=torch.int64), multimodal_params=[{} for _ in range(rows)], - b_logits_rows=accepted, + b_logits_row_idx=accepted, ) @@ -113,9 +113,9 @@ def test_first_draft_selects_accepted_tail_across_padding(draft_model, eager, re output = model._decode(model_input) torch.testing.assert_close(output.logits[:, 0], accepted.float()) torch.testing.assert_close(output.mtp_collector.spec_hidden[:, 0], model_input.input_ids.float()) - assert output.logits_row_stride == 5 + assert output.input_rows_per_logit == 5 assert model_input.batch_size == requests * 5 - assert model_input.b_logits_rows is accepted + assert model_input.b_logits_row_idx is accepted if not eager and requests == 52: assert 320 in model.draft_tail_graph.graph assert not model.graph.graph @@ -130,7 +130,7 @@ def test_later_draft_keeps_one_output_per_request(draft_model, rows): model_input = make_input(rows) model_input.input_ids += step output = model._decode(model_input) - assert output.logits_row_stride == 1 + assert output.input_rows_per_logit == 1 torch.testing.assert_close(output.logits[:, 0], model_input.input_ids.float()) assert output.mtp_collector.spec_hidden.shape == (rows, 1) if rows == 160: @@ -162,6 +162,6 @@ def allocate(*args, **kwargs): forward=lambda model_input: seen.append(model_input), ) graph = cuda_graph.CudaGraph(5, 5, 5, max_batch_size=10) - graph.warmup(model, logits_row_stride=5) + graph.warmup(model, input_rows_per_logit=5) assert [x.batch_size for x in seen] == [10, 5] - assert [x.b_logits_rows.tolist() for x in seen] == [[4, 9], [4]] + assert [x.b_logits_row_idx.tolist() for x in seen] == [[4, 9], [4]] diff --git a/unit_tests/common/basemodel/test_model_output.py b/unit_tests/common/basemodel/test_model_output.py index d06997adf4..7a9045a344 100644 --- a/unit_tests/common/basemodel/test_model_output.py +++ b/unit_tests/common/basemodel/test_model_output.py @@ -333,14 +333,14 @@ def test_selected_draft_logits_unpad_preserves_physical_hidden_rows(origin): logits=torch.arange(12).view(3, 4), logits_token_ids=torch.arange(12).view(3, 4), mtp_collector=ModelMtpOutputCollector(spec_hidden=torch.arange(45).view(15, 3)), - logits_row_stride=5, + input_rows_per_logit=5, ) unpadded = model._create_unpad_decode_model_output(output, origin) assert unpadded.logits.shape == (origin // 5, 4) assert unpadded.logits_token_ids.shape == (origin // 5, 4) assert unpadded.mtp_collector.spec_hidden.shape == (origin, 3) assert output.mtp_collector.spec_hidden.shape == (15, 3) - assert unpadded.logits_row_stride == 5 + assert unpadded.input_rows_per_logit == 5 @pytest.mark.parametrize("disabled", [None, "dynamic", "overlap", "tp", "projection", "model", "target"]) @@ -358,7 +358,7 @@ def test_draft_tail_scope(disabled): model.args.enable_decode_microbatch_overlap = disabled == "overlap" if disabled == "projection": model.args.ssm_replay_config = {"compute_mode": "inline"} - assert model.use_replayssm_draft_tail() == (disabled is None) + assert model.should_project_draft_tail_logits() == (disabled is None) def _create_empty_decode_input(): diff --git a/unit_tests/common/basemodel/triton_kernel/linear_att/test_acceptance_fusion.py b/unit_tests/common/basemodel/triton_kernel/linear_att/test_acceptance_fusion.py index f431878851..483f8d50d9 100644 --- a/unit_tests/common/basemodel/triton_kernel/linear_att/test_acceptance_fusion.py +++ b/unit_tests/common/basemodel/triton_kernel/linear_att/test_acceptance_fusion.py @@ -3,8 +3,8 @@ import triton from lightllm.common.basemodel.triton_kernel.mtp_utils import linear_att_mtp_state_index_update as update -from lightllm.common.basemodel.triton_kernel.linear_att.replayssm import ReplaySSMCache -from lightllm.common.recurrent_state import RecurrentState +from lightllm.common.basemodel.triton_kernel.linear_att.replayssm import ReplaySSMBackend +from lightllm.common.linear_state import LinearStateManager from lightllm.common.req_manager.linear_att import ReqManagerForMamba @@ -21,24 +21,26 @@ def make(n, width): mtps.extend(range(l)) tensor = lambda x: torch.tensor(x, device="cuda", dtype=torch.int32) args = [tensor(starts), tensor(reqs), tensor(mtps), tensor([1] * len(reqs)), width] - cache = ReplaySSMCache.__new__(ReplaySSMCache) - cache.cursors = tensor([7] * (hold + 1)) - cache.hold = hold - cache.capacity = max(4, triton.next_power_of_2(width)) - cache.verify_width = width + ssm_backend = ReplaySSMBackend.__new__(ReplaySSMBackend) + ssm_backend.cursors = tensor([7] * (hold + 1)) + ssm_backend.hold_req_id = hold + ssm_backend.capacity = max(4, triton.next_power_of_2(width)) + ssm_backend.verify_width = width manager = ReqManagerForMamba.__new__(ReqManagerForMamba) - manager.recurrent_state = object.__new__(RecurrentState) - manager.recurrent_state.accepted_index = tensor([-9] * (hold + 1)) - manager.recurrent_state.updates = cache + manager.linear_state_manager = object.__new__(LinearStateManager) + manager.linear_state_manager.accepted_state_index = tensor([-9] * (hold + 1)) + manager.linear_state_manager.ssm_backend = ssm_backend return manager, args, (ids, lengths, starts) -def old(m, args): - update(m.recurrent_state.accepted_index, *args) - m.recurrent_state.updates.accept_updates(args[1][args[0].long()], m.recurrent_state.accepted_index) +def accept_unfused(m, args): + update(m.linear_state_manager.accepted_state_index, *args) + m.linear_state_manager.ssm_backend.accept_updates( + args[1][args[0].long()], m.linear_state_manager.accepted_state_index + ) -def new(m, args): +def accept_fused(m, args): m.update_mtp_state(*args) @@ -50,80 +52,84 @@ def test_acceptance_fusion_matches_original_and_cpu_under_graph(): m, args, (ids, lengths, starts) = make(n, width) ref, _, _ = make(n, width) # Compile both paths before graph capture, then reset before replay. - old(ref, args) - new(m, args) + accept_unfused(ref, args) + accept_fused(m, args) torch.cuda.synchronize() graph = torch.cuda.CUDAGraph() with torch.cuda.graph(graph): - new(m, args) - expected_index = [-9] * len(m.recurrent_state.accepted_index) + accept_fused(m, args) + expected_index = [-9] * len(m.linear_state_manager.accepted_state_index) expected_cursor = [7] * len(expected_index) - m.recurrent_state.accepted_index.fill_(-9) - ref.recurrent_state.accepted_index.fill_(-9) - m.recurrent_state.updates.cursors.fill_(7) - ref.recurrent_state.updates.cursors.fill_(7) + m.linear_state_manager.accepted_state_index.fill_(-9) + ref.linear_state_manager.accepted_state_index.fill_(-9) + m.linear_state_manager.ssm_backend.cursors.fill_(7) + ref.linear_state_manager.ssm_backend.cursors.fill_(7) for step in range(24): mask = [] for i, (r, l) in enumerate(zip(ids, lengths)): count = (step + i) % (l + 1) mask.extend([1] * count + [0] * (l - count)) expected_index[r] = count - 1 - if r != m.recurrent_state.updates.hold: + if r != m.linear_state_manager.ssm_backend.hold_req_id: expected_cursor[r] += count if width > 1 else 1 args[3].copy_(torch.tensor(mask, device="cuda", dtype=torch.int32)) - old(ref, args) + accept_unfused(ref, args) if step % 2: graph.replay() else: - new(m, args) - assert torch.equal(m.recurrent_state.accepted_index, ref.recurrent_state.accepted_index) - assert torch.equal(m.recurrent_state.updates.cursors, ref.recurrent_state.updates.cursors) - assert m.recurrent_state.accepted_index.tolist() == expected_index - assert m.recurrent_state.updates.cursors.tolist() == expected_cursor + accept_fused(m, args) + assert torch.equal( + m.linear_state_manager.accepted_state_index, ref.linear_state_manager.accepted_state_index + ) + assert torch.equal( + m.linear_state_manager.ssm_backend.cursors, ref.linear_state_manager.ssm_backend.cursors + ) + assert m.linear_state_manager.accepted_state_index.tolist() == expected_index + assert m.linear_state_manager.ssm_backend.cursors.tolist() == expected_cursor checks += 1 - # Native optional-pointer path and manager dispatch leave metadata-only cache absent. - m.recurrent_state.updates = None - new(m, args) - assert m.recurrent_state.accepted_index.tolist() == expected_index + # Native optional-pointer path and manager dispatch leave metadata-only backend absent. + m.linear_state_manager.ssm_backend = None + accept_fused(m, args) + assert m.linear_state_manager.accepted_state_index.tolist() == expected_index - # Non-Replay caches must still receive the accepted request IDs and indexes. + # Non-Replay backends must still receive the accepted request IDs and indexes. class CompactProbe: def accept_updates(self, reqs, index): self.seen = (reqs.clone(), index.clone()) m, args, (ids, _, _) = make(7, 4) probe = CompactProbe() - m.recurrent_state.updates = probe - new(m, args) + m.linear_state_manager.ssm_backend = probe + accept_fused(m, args) assert probe.seen[0].tolist() == ids - assert torch.equal(probe.seen[1], m.recurrent_state.accepted_index) + assert torch.equal(probe.seen[1], m.linear_state_manager.accepted_state_index) print("ACCEPT_EQUIVALENCE_CHECKS=" + str(checks), flush=True) @pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") @pytest.mark.parametrize("dtype", [torch.float32, torch.bfloat16]) -@pytest.mark.parametrize("projection_mode", ["inline", "precompute"]) -def test_fused_acceptance_preserves_folds_snapshots_and_slot_reuse(dtype, projection_mode): +@pytest.mark.parametrize("compute_mode", ["inline", "precompute"]) +def test_fused_acceptance_preserves_folds_snapshots_and_slot_reuse(dtype, compute_mode): torch.manual_seed(93) initial = torch.randn(2, 8, 4, 32, 32, device="cuda").to(dtype) * 0.01 - caches = [ - ReplaySSMCache( + ssm_backends = [ + ReplaySSMBackend( initial.clone(), 8, 4, torch.bfloat16, num_key_heads=2, - projection_mode=projection_mode, + compute_mode=compute_mode, run_config={"BV": 32, "num_warps": 1, "num_stages": 1}, ) for _ in range(2) ] managers = [] - for cache in caches: + for ssm_backend in ssm_backends: manager = ReqManagerForMamba.__new__(ReqManagerForMamba) - manager.recurrent_state = object.__new__(RecurrentState) - manager.recurrent_state.updates = cache - manager.recurrent_state.accepted_index = torch.zeros(8, dtype=torch.int32, device="cuda") + manager.linear_state_manager = object.__new__(LinearStateManager) + manager.linear_state_manager.ssm_backend = ssm_backend + manager.linear_state_manager.accepted_state_index = torch.zeros(8, dtype=torch.int32, device="cuda") managers.append(manager) reqs = torch.full((4,), 7, dtype=torch.int32, device="cuda") cu = torch.tensor([0, 3, 6, 9, 12], dtype=torch.int32, device="cuda") @@ -139,10 +145,10 @@ def test_fused_acceptance_preserves_folds_snapshots_and_slot_reuse(dtype, projec args = [cu[:-1], token_reqs, mtp_indexes, accepted_mask, 4] def cycle(i): - cache = caches[i] - cache.prepare_decode(reqs, cu) - outputs = [cache.forward(layer, q, k, v, a, b, alog, alog, reqs, cu) for layer in range(2)] - (old if i == 0 else new)(managers[i], args) + ssm_backend = ssm_backends[i] + ssm_backend.prepare_decode(reqs, cu) + outputs = [ssm_backend.forward(layer, q, k, v, a, b, alog, alog, reqs, cu)[0] for layer in range(2)] + (accept_unfused if i == 0 else accept_fused)(managers[i], args) return outputs graphs, outputs = [], [] @@ -180,21 +186,23 @@ def cycle(i): graph.replay() for left, right in zip(outputs[0], outputs[1]): assert torch.equal(left, right) - assert torch.equal(caches[0].state, caches[1].state) - assert torch.equal(caches[0].cursors, caches[1].cursors) - assert torch.equal(managers[0].recurrent_state.accepted_index, managers[1].recurrent_state.accepted_index) + assert torch.equal(ssm_backends[0].state, ssm_backends[1].state) + assert torch.equal(ssm_backends[0].cursors, ssm_backends[1].cursors) + assert torch.equal( + managers[0].linear_state_manager.accepted_state_index, managers[1].linear_state_manager.accepted_state_index + ) if iteration % 7 == 0: - assert torch.equal(caches[0].snapshot_accepted_state(3), caches[1].snapshot_accepted_state(3)) + assert torch.equal(ssm_backends[0].get_accepted_state(3), ssm_backends[1].get_accepted_state(3)) if iteration % 11 == 10: - for cache in caches: - cache.merge_accepted_updates(reqs[:3]) - assert torch.equal(caches[0].state, caches[1].state) + for ssm_backend in ssm_backends: + ssm_backend.merge_accepted_updates(reqs[:3]) + assert torch.equal(ssm_backends[0].state, ssm_backends[1].state) if iteration == 19: for manager in managers: - manager.recurrent_state.updates.clear_history(1) - manager.recurrent_state.updates.state[:, 1].zero_() - manager.recurrent_state.accepted_index[1] = 0 - for cache in caches: - cache.merge_accepted_updates(reqs[:3]) - assert torch.equal(caches[0].state, caches[1].state) - assert torch.equal(caches[0].state[:, [0, 2, 4, 5, 7]], initial[:, [0, 2, 4, 5, 7]]) + manager.linear_state_manager.ssm_backend.clear_history(1) + manager.linear_state_manager.ssm_backend.state[:, 1].zero_() + manager.linear_state_manager.accepted_state_index[1] = 0 + for ssm_backend in ssm_backends: + ssm_backend.merge_accepted_updates(reqs[:3]) + assert torch.equal(ssm_backends[0].state, ssm_backends[1].state) + assert torch.equal(ssm_backends[0].state[:, [0, 2, 4, 5, 7]], initial[:, [0, 2, 4, 5, 7]]) diff --git a/unit_tests/common/basemodel/triton_kernel/linear_att/test_mtp_state_params.py b/unit_tests/common/basemodel/triton_kernel/linear_att/test_mtp_state_params.py index 6de439d7b5..3c5ef837d3 100644 --- a/unit_tests/common/basemodel/triton_kernel/linear_att/test_mtp_state_params.py +++ b/unit_tests/common/basemodel/triton_kernel/linear_att/test_mtp_state_params.py @@ -1,6 +1,8 @@ import pytest import torch +from lightllm.common.linear_state import LinearStateManager + from lightllm.common.basemodel.triton_kernel.linear_att.mtp_state_params import ( build_dynamic_mtp_linear_att_state_params, ) @@ -12,17 +14,22 @@ def test_dynamic_sequences_stay_within_request_capacity_under_graph(request_capa hold = request_capacity reqs = torch.full((token_bucket,), hold, dtype=torch.int32, device="cuda") mtp_index = torch.zeros_like(reqs) - accepted = torch.arange(request_capacity + 1, dtype=torch.int32, device="cuda") % 4 + manager = object.__new__(LinearStateManager) + manager.accepted_state_index = torch.arange(request_capacity + 1, dtype=torch.int32, device="cuda") % 4 + manager.ssm_slots_per_req = 4 + manager.ssm_backend = None sequences = min(request_capacity, token_bucket) def build(): - return build_dynamic_mtp_linear_att_state_params(reqs, mtp_index, accepted, hold) + cu, ids = build_dynamic_mtp_linear_att_state_params(reqs, mtp_index, hold) + ids, ssm, counts = manager.prepare_mtp_decode(ids, cu) + return cu, ids, ssm, counts build() torch.cuda.synchronize() graph = torch.cuda.CUDAGraph() with torch.cuda.graph(graph): - cu, ids, counts = build() + cu, ids, ssm, counts = build() assert cu.numel() == sequences + 1 assert ids.numel() == counts.numel() == sequences @@ -46,7 +53,8 @@ def build(): expected_cu.append(expected_cu[-1] + length) expected_cu += [tokens] * (sequences + 1 - len(expected_cu)) expected_ids = request_ids + [hold] * (sequences - len(request_ids)) - expected_counts = [req % 4 + 1 for req in request_ids] + [1] * (sequences - len(request_ids)) assert cu.tolist() == expected_cu assert ids.tolist() == expected_ids + expected_counts = [req % 4 + 1 for req in request_ids] + [1] * (sequences - len(request_ids)) assert counts.tolist() == expected_counts + assert ssm.tolist() == [[req * 4 + offset for offset in range(4)] for req in expected_ids] 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 index 5bcacc8ecf..ebbb915a15 100644 --- a/unit_tests/common/basemodel/triton_kernel/linear_att/test_replayssm.py +++ b/unit_tests/common/basemodel/triton_kernel/linear_att/test_replayssm.py @@ -1,20 +1,20 @@ import pytest import torch -from lightllm.common.basemodel.triton_kernel.linear_att.replayssm import ReplaySSMCache +from lightllm.common.basemodel.triton_kernel.linear_att.replayssm import ReplaySSMBackend from lightllm.common.basemodel.triton_kernel.linear_att.ssm_autotune import get_configs pytestmark = pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") -def reference(q, k, v, a, b, a_log, bias, state, kda=False, lower_bound=-5.0): +def reference(q, k, v, a, b, a_log, bias, state, kda=False, gate_lower_bound=-5.0): 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) if kda: - g = lower_bound * torch.sigmoid(a_log.exp()[:, None] * (a.float().view_as(k) + bias.view_as(k))) + g = gate_lower_bound * torch.sigmoid(a_log.exp()[:, None] * (a.float().view_as(k) + bias.view_as(k))) state = state * g.exp()[..., None] else: g = -a_log.exp() * torch.nn.functional.softplus(a.float() + bias) @@ -30,14 +30,14 @@ def reference(q, k, v, a, b, a_log, bias, state, kda=False, lower_bound=-5.0): @pytest.mark.parametrize("kda", [False, True], ids=["gdn", "kda"]) def test_fp16_correction_history_preserves_finite_recurrence(capacity, bv, state_dtype, kda): state = torch.zeros(1, 2, 1, 128, 128, dtype=state_dtype, device="cuda") - cache = ReplaySSMCache( + ssm_backend = ReplaySSMBackend( state, capacity, 2, torch.float16, num_key_heads=1, kda=kda, - projection_mode="precompute", + compute_mode="precompute", run_config={"BV": bv, "num_warps": 1}, ) reqs = torch.tensor([0, 1], dtype=torch.int32, device="cuda") @@ -56,13 +56,13 @@ def test_fp16_correction_history_preserves_finite_recurrence(capacity, bv, state bias = torch.zeros(128 if kda else 1, device="cuda") expected_state = state[0, 0].float().clone() for round_idx in range(2): - cache.prepare_decode(reqs, cu) + ssm_backend.prepare_decode(reqs, cu) # The second round must consume the first round's accepted history. # For capacity 4 the next span exactly fills it, so no fold should hide # an overflowed derived correction by reconstructing from raw inputs. - assert cache.cursors[0].item() == 2 * round_idx + assert ssm_backend.cursors[0].item() == 2 * round_idx assert not state.any() - actual = cache.forward(0, q, k, v, a, b, alog, bias, reqs, cu) + actual = ssm_backend.forward(0, q, k, v, a, b, alog, bias, reqs, cu)[0] expected = [] for token in range(2): output, expected_state = reference( @@ -72,13 +72,13 @@ def test_fp16_correction_history_preserves_finite_recurrence(capacity, bv, state assert torch.isfinite(expected_state).all() assert torch.isfinite(actual).all() torch.testing.assert_close(actual[0, :2].float(), torch.stack(expected), atol=4, rtol=1e-3) - cache.accept_updates(reqs, accepted) + ssm_backend.accept_updates(reqs, accepted) # The correction is larger than FP16 can represent even though the # recurrent state, inputs, and attention outputs are all finite. - assert cache.deltas[0, 0, 0, 1, 0].item() > torch.finfo(torch.float16).max - assert torch.isfinite(cache.deltas[0, 0, :, :4]).all() - assert cache.cursors[1].item() == 0 - torch.testing.assert_close(cache.snapshot_accepted_state(0)[0].float(), expected_state, atol=128, rtol=0.01) + assert ssm_backend.deltas[0, 0, 0, 1, 0].item() > torch.finfo(torch.float16).max + assert torch.isfinite(ssm_backend.deltas[0, 0, :, :4]).all() + assert ssm_backend.cursors[1].item() == 0 + torch.testing.assert_close(ssm_backend.get_accepted_state(0)[0].float(), expected_state, atol=128, rtol=0.01) @pytest.mark.parametrize("capacity", [16, 32, 64]) @@ -91,26 +91,26 @@ def test_replay_acceptance_flush_and_graph( capacity, dtype, run_config=None, - projection_mode="inline", + compute_mode="inline", activation_dtype=torch.bfloat16, gate_shift=-3, head_geometry=(2, 4), kda=False, - lower_bound=-5.0, + gate_lower_bound=-5.0, ): torch.manual_seed(21) kdim, vdim = dims h, hv = head_geometry slots = 5 state = (torch.randn(2, slots, hv, kdim, vdim, device="cuda") * 0.01).to(dtype) - cache = ReplaySSMCache( + ssm_backend = ReplaySSMBackend( state, capacity, width, num_key_heads=h, kda=kda, - lower_bound=lower_bound, - projection_mode=projection_mode, + gate_lower_bound=gate_lower_bound, + compute_mode=compute_mode, run_config=run_config, activation_dtype=activation_dtype, ) @@ -135,16 +135,16 @@ def test_replay_acceptance_flush_and_graph( accepted = torch.zeros(slots, dtype=torch.int32, device="cuda") def step(): - cache.prepare_decode(reqs, cu) - out = [cache.forward(layer, q, k, v, a, b, alog, bias, reqs, cu) for layer in range(2)] - cache.accept_updates(reqs, accepted) + ssm_backend.prepare_decode(reqs, cu) + out = [ssm_backend.forward(layer, q, k, v, a, b, alog, bias, reqs, cu)[0] for layer in range(2)] + ssm_backend.accept_updates(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_() + ssm_backend.cursors.zero_() reqs.fill_(slots - 1) torch.cuda.synchronize() graph = torch.cuda.CUDAGraph() @@ -169,17 +169,26 @@ def step(): 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, kda=kda, lower_bound=lower_bound + q[0, t], + k[0, t], + v[0, t], + a[t], + b[t], + alog, + bias, + cur, + kda=kda, + gate_lower_bound=gate_lower_bound, ) 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.merge_accepted_updates(reqs[:2]) + ssm_backend.merge_accepted_updates(reqs[:2]) expected = expected.to(dtype).float() pending = [0, 0] torch.testing.assert_close(state[:, :4].float(), expected[:, :4], atol=state_atol, rtol=state_rtol) - cache.merge_accepted_updates(reqs[:2]) + ssm_backend.merge_accepted_updates(reqs[:2]) torch.testing.assert_close(state.float(), expected.to(dtype).float(), atol=state_atol, rtol=state_rtol) @@ -187,14 +196,14 @@ def step(): "state_dtype,activation_dtype", [(torch.float32, torch.float32), (torch.float32, torch.bfloat16), (torch.bfloat16, torch.bfloat16)], ) -@pytest.mark.parametrize("projection_mode", ["inline", "precompute"]) +@pytest.mark.parametrize("compute_mode", ["inline", "precompute"]) @pytest.mark.parametrize("dims", [(32, 64), (128, 128)]) @pytest.mark.parametrize("width", [1, 4]) @pytest.mark.parametrize("gate_shift", [-3, -10]) def test_fold_matches_reference_across_precisions( state_dtype, activation_dtype, - projection_mode, + compute_mode, dims, width, gate_shift, @@ -205,7 +214,7 @@ def test_fold_matches_reference_across_precisions( 8, state_dtype, run_config={"BV": 32, "num_warps": 1, "num_stages": 3}, - projection_mode=projection_mode, + compute_mode=compute_mode, activation_dtype=activation_dtype, gate_shift=gate_shift, ) @@ -219,14 +228,14 @@ def test_short_replay_history_matches_independent_reference(capacity, width, dty @pytest.mark.parametrize("capacity,width", [(4, 1), (4, 3), (4, 4), (8, 1), (8, 4), (8, 5), (16, 4), (16, 16), (32, 4)]) @pytest.mark.parametrize("dtype", [torch.float32, torch.bfloat16]) -def test_precomputed_state_matches_independent_reference(capacity, width, dtype): +def test_precomputed_projections_match_independent_reference(capacity, width, dtype): test_replay_acceptance_flush_and_graph( width, (128, 128), capacity, dtype, run_config={"BV": 32, "num_warps": 1}, - projection_mode="precompute", + compute_mode="precompute", ) @@ -238,7 +247,7 @@ def test_production_gdn_geometry_matches_independent_reference(dtype): 8, dtype, run_config={"BV": 32, "num_warps": 1, "num_stages": 3}, - projection_mode="precompute", + compute_mode="precompute", head_geometry=(4, 12), ) @@ -251,18 +260,18 @@ def test_production_gdn_geometry_matches_independent_reference(dtype): ) @pytest.mark.parametrize("bv", [8, 32]) @pytest.mark.parametrize("width", [1, 4]) -def test_precomputed_state_preserves_projection_precision(warps, decay, state_dtype, dtype, width, bv): +def test_precomputed_projections_preserve_precision(warps, decay, state_dtype, dtype, width, bv): torch.manual_seed(35) state = (torch.randn((1, 3, 4, 128, 128), device="cuda") * 0.1).to(state_dtype) config = {"BV": bv, "num_warps": warps} - caches = [ - ReplaySSMCache( + ssm_backends = [ + ReplaySSMBackend( state.clone(), 16 if width == 1 else 8, width, dtype, num_key_heads=2, - projection_mode=mode, + compute_mode=mode, run_config=config, ) for mode in ("inline", "precompute") @@ -283,10 +292,10 @@ def test_precomputed_state_preserves_projection_precision(warps, decay, state_dt accepted[0] = step % max(1, width - 1) accepted[1] = step % width outputs = [] - for cache in caches: - cache.prepare_decode(reqs, cu) - outputs.append(cache.forward(0, q, k, v, a, b, alog, alog, reqs, cu)) - cache.accept_updates(reqs, accepted) + for ssm_backend in ssm_backends: + ssm_backend.prepare_decode(reqs, cu) + outputs.append(ssm_backend.forward(0, q, k, v, a, b, alog, alog, reqs, cu)[0]) + ssm_backend.accept_updates(reqs, accepted) if dtype == torch.float32: torch.testing.assert_close(outputs[0], outputs[1], atol=2e-6, rtol=2e-5) else: @@ -294,24 +303,24 @@ def test_precomputed_state_preserves_projection_precision(warps, decay, state_dt error = torch.linalg.vector_norm(outputs[0].float() - outputs[1].float()) scale = torch.linalg.vector_norm(outputs[0].float()) assert error < scale * (torch.finfo(dtype).eps / 2) - for cache in caches: - cache.merge_accepted_updates(reqs) - torch.testing.assert_close(caches[0].state, caches[1].state, atol=0, rtol=0) + for ssm_backend in ssm_backends: + ssm_backend.merge_accepted_updates(reqs) + torch.testing.assert_close(ssm_backends[0].state, ssm_backends[1].state, atol=0, rtol=0) def test_bf16_decode_fold_uses_rounded_checkpoint(): # Automatic fold and explicit materialization must start from the same BF16 # checkpoint, including the token computed inside the automatic-fold kernel. state = torch.full((1, 2, 1, 32, 64), 0.123, device="cuda", dtype=torch.bfloat16) - caches = [ReplaySSMCache(state.clone(), 16, 1, torch.float32) for _ in range(2)] + ssm_backends = [ReplaySSMBackend(state.clone(), 16, 1, torch.float32) for _ in range(2)] reqs = torch.tensor([0], device="cuda", dtype=torch.int32) - for cache in caches: - cache.raw_keys.zero_() - cache.raw_values.zero_() - cache.gates.fill_(-0.01) - cache.betas.zero_() - cache.cursors[0] = 16 - caches[1].merge_accepted_updates(reqs) + for ssm_backend in ssm_backends: + ssm_backend.raw_keys.zero_() + ssm_backend.raw_values.zero_() + ssm_backend.gates.fill_(-0.01) + ssm_backend.betas.zero_() + ssm_backend.cursors[0] = 16 + ssm_backends[1].merge_accepted_updates(reqs) q = torch.ones((1, 1, 32), device="cuda") k = torch.zeros_like(q) v = torch.zeros((1, 1, 64), device="cuda") @@ -320,22 +329,22 @@ def test_bf16_decode_fold_uses_rounded_checkpoint(): alog = torch.zeros(1, device="cuda") bias = torch.zeros_like(alog) outputs = [] - for cache in caches: - cache.prepare_decode(reqs) - cache.accept_updates(reqs) - outputs.append(cache.forward(0, q, k, v, a, b, alog, bias, reqs)) + for ssm_backend in ssm_backends: + ssm_backend.prepare_decode(reqs) + ssm_backend.accept_updates(reqs) + outputs.append(ssm_backend.forward(0, q, k, v, a, b, alog, bias, reqs)[0]) torch.testing.assert_close(outputs[0], outputs[1], atol=1e-7, rtol=1e-6) - torch.testing.assert_close(caches[0].state, caches[1].state, atol=0, rtol=0) + torch.testing.assert_close(ssm_backends[0].state, ssm_backends[1].state, atol=0, rtol=0) @pytest.mark.parametrize("dtype", [torch.float32, torch.bfloat16]) -@pytest.mark.parametrize("projection_mode", ["inline", "precompute"]) +@pytest.mark.parametrize("compute_mode", ["inline", "precompute"]) @pytest.mark.parametrize("capture_metadata", [False, True]) -def test_dynamic_verify_lengths_fill_history_before_folding(dtype, projection_mode, capture_metadata): +def test_dynamic_verify_lengths_fill_history_before_folding(dtype, compute_mode, capture_metadata): torch.manual_seed(31) state = (torch.randn(2, 4, 4, 128, 128, device="cuda") * 0.01).to(dtype) - cache = ReplaySSMCache( - state, 8, 4, dtype, num_key_heads=2, projection_mode=projection_mode, run_config={"BV": 32, "num_warps": 1} + ssm_backend = ReplaySSMBackend( + state, 8, 4, dtype, num_key_heads=2, compute_mode=compute_mode, run_config={"BV": 32, "num_warps": 1} ) expected = state.float().clone() pending, phases = [0, 0], [0, 0] @@ -350,12 +359,12 @@ def test_dynamic_verify_lengths_fill_history_before_folding(dtype, projection_mo alog = torch.zeros(4, device="cuda") def forward(): - return [cache.forward(layer, q, k, v, a, b, alog, alog, reqs, cu) for layer in range(2)] + return [ssm_backend.forward(layer, q, k, v, a, b, alog, alog, reqs, cu)[0] for layer in range(2)] def step(): - positions = cache.prepare_decode(reqs, cu) + positions = ssm_backend.prepare_decode(reqs, cu) out = forward() - cache.accept_updates(reqs, accepted) + ssm_backend.accept_updates(reqs, accepted) return positions, out step() @@ -389,10 +398,10 @@ def step(): pending[req] = 0 phases[req] ^= 16 if not capture_metadata: - positions.copy_(cache.prepare_decode(reqs, cu)) + positions.copy_(ssm_backend.prepare_decode(reqs, cu)) graph.replay() if not capture_metadata: - cache.accept_updates(reqs, accepted) + ssm_backend.accept_updates(reqs, accepted) assert positions.tolist() == old_cursors for seq, (req, width, count) in enumerate(zip(ids, widths, counts)): for layer in range(2): @@ -404,14 +413,14 @@ def step(): if j + 1 == count: expected[layer, req].copy_(current) pending[req] += count - assert cache.cursors.tolist() == [phases[i] + pending[i] for i in range(2)] + [0, 0] + assert ssm_backend.cursors.tolist() == [phases[i] + pending[i] for i in range(2)] + [0, 0] for output in outputs: assert torch.count_nonzero(output[:, sum(widths) :]).item() == 0 torch.testing.assert_close(state[:, 2:], before[:, 2:], atol=0, rtol=0) if iteration == 2: assert pending[0] == 8 and phases[0] == 0 torch.testing.assert_close(state[:, 0], before[:, 0], atol=0, rtol=0) - cache.merge_accepted_updates(reqs[:2]) + ssm_backend.merge_accepted_updates(reqs[:2]) torch.testing.assert_close( state.float(), expected.to(dtype).float(), @@ -422,20 +431,20 @@ def step(): def test_snapshot_without_history_copies_active_state(): state = torch.randn((2, 3, 2, 32, 64), device="cuda", dtype=torch.bfloat16) - cache = ReplaySSMCache(state, 4, 3) - snapshot = cache.snapshot_accepted_state(1) + ssm_backend = ReplaySSMBackend(state, 4, 3) + snapshot = ssm_backend.get_accepted_state(1) torch.testing.assert_close(snapshot, state[:, 1], rtol=0, atol=0) - assert cache.cursors[1].item() == 0 + assert ssm_backend.cursors[1].item() == 0 @pytest.mark.parametrize("dtype", [torch.float32, torch.bfloat16]) def test_compact_preserves_accepted_prefix(dtype, run_config=None, kda=False): - from lightllm.common.basemodel.triton_kernel.linear_att.replayssm_compact import CompactSSMCache + from lightllm.common.basemodel.triton_kernel.linear_att.compact_ssm import CompactSSMBackend torch.manual_seed(17) hv, kd, vd, width = 4, 128 if kda else 64, 128 if kda else 64, 3 state = torch.randn(2, 4, hv, kd, vd, device="cuda", dtype=dtype) * 0.01 - cache = CompactSSMCache(state, width, torch.bfloat16, num_key_heads=hv, kda=kda, run_config=run_config) + ssm_backend = CompactSSMBackend(state, width, torch.bfloat16, num_key_heads=hv, kda=kda, run_config=run_config) 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) @@ -448,9 +457,9 @@ def test_compact_preserves_accepted_prefix(dtype, run_config=None, kda=False): bias = torch.randn(hv * kd if 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, cu) for layer in range(2)] + outputs = [ssm_backend.forward(layer, q, k, v, a, b, alog, bias, reqs, cu)[0] for layer in range(2)] assert torch.equal(state, before), "verify must not mutate the checkpoint" - cache.accept_updates(reqs, accepted) + ssm_backend.accept_updates(reqs, accepted) for layer in range(2): for req, start, length in [(2, 0, 3), (0, 3, 2)]: cur = before[layer, req].float() @@ -470,7 +479,7 @@ def test_compact_preserves_accepted_prefix(dtype, run_config=None, kda=False): @pytest.mark.parametrize("dtype", [torch.float32, torch.bfloat16]) @pytest.mark.parametrize("width", [3, 4]) def test_compact_matches_native_mtp(dtype, width, run_config=None): - from lightllm.common.basemodel.triton_kernel.linear_att.replayssm_compact import CompactSSMCache + from lightllm.common.basemodel.triton_kernel.linear_att.compact_ssm import CompactSSMBackend from lightllm.common.basemodel.triton_kernel.linear_att.mtp_fused_recurrent import ( mtp_fused_recurrent_gated_delta_rule, ) @@ -491,14 +500,14 @@ def test_compact_matches_native_mtp(dtype, width, run_config=None): 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, num_key_heads=h, run_config=run_config) + ssm_backend = CompactSSMBackend(state, width, torch.bfloat16, num_key_heads=h, run_config=run_config) 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, cu) + out = ssm_backend.forward(0, q, k, v, a, b, alog, bias, reqs, cu)[0] ref, _ = mtp_fused_recurrent_gated_delta_rule( q, k, @@ -514,7 +523,7 @@ def test_compact_matches_native_mtp(dtype, width, run_config=None): b, run_config={"num_stages": 1, **(run_config or {"BV": 8, "num_warps": 1})}, ) - cache.accept_updates(reqs, accepted) + ssm_backend.accept_updates(reqs, accepted) # BF16 casts can straddle a rounding midpoint across compiler/layouts. rtol = 2 * torch.finfo(dtype).eps if dtype == torch.bfloat16 else 0 torch.testing.assert_close(out, ref, rtol=rtol, atol=2e-5 if dtype == torch.bfloat16 else 0) @@ -542,7 +551,7 @@ def test_autotune_candidates_match_recurrence(dtype, run_config, mode): @pytest.mark.parametrize("dtype", [torch.float32, torch.bfloat16]) def test_checkpoint_restores_pending_history_and_accepted_conv(save_big, monkeypatch, dtype): from types import SimpleNamespace - from lightllm.common.recurrent_state import RecurrentState + from lightllm.common.linear_state import LinearStateManager 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 @@ -551,7 +560,7 @@ def test_checkpoint_restores_pending_history_and_accepted_conv(save_big, monkeyp state_rtol = 2e-5 if dtype == torch.float32 else torch.finfo(dtype).eps # Real save/restore and PD helpers, with just the unrelated allocator omitted. manager = object.__new__(ReqManagerForMamba) - manager.recurrent_state = object.__new__(RecurrentState) + manager.linear_state_manager = object.__new__(LinearStateManager) config = SimpleNamespace( linear_layer_num=2, conv_state_dtype=torch.bfloat16, @@ -568,38 +577,40 @@ def test_checkpoint_restores_pending_history_and_accepted_conv(save_big, monkeyp conv_kernel_size=4, ) manager.linear_config = config - manager.recurrent_state.linear_config = config + manager.linear_state_manager.linear_config = config manager.mtp_step = 2 - manager.recurrent_state.slots_per_req = 1 - manager.recurrent_state.accepted_index = torch.zeros(4, device="cuda", dtype=torch.int32) - manager.recurrent_state.accepted_index[1] = 1 - manager.recurrent_state.conv_state = LayerCache(4, torch.bfloat16, (192, 5), 2, "cuda") - manager.recurrent_state.conv_state.buffer.copy_(torch.randn_like(manager.recurrent_state.conv_state.buffer)) - manager.recurrent_state.ssm_state = LayerCache(4, dtype, (2, 32, 64), 2, "cuda") - manager.recurrent_state.ssm_state.buffer.zero_() - manager.recurrent_state.updates = ReplaySSMCache(manager.recurrent_state.ssm_state.buffer, 16, 3) - cache = manager.recurrent_state.updates - 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 - active_before = manager.recurrent_state.ssm_state.buffer[:, 1].clone() - expected = manager.recurrent_state.ssm_state.buffer[:, 1].float().clone() + manager.linear_state_manager.ssm_slots_per_req = 1 + manager.linear_state_manager.accepted_state_index = torch.zeros(4, device="cuda", dtype=torch.int32) + manager.linear_state_manager.accepted_state_index[1] = 1 + manager.linear_state_manager.conv_state = LayerCache(4, torch.bfloat16, (192, 5), 2, "cuda") + manager.linear_state_manager.conv_state.buffer.copy_( + torch.randn_like(manager.linear_state_manager.conv_state.buffer) + ) + manager.linear_state_manager.ssm_state = LayerCache(4, dtype, (2, 32, 64), 2, "cuda") + manager.linear_state_manager.ssm_state.buffer.zero_() + manager.linear_state_manager.ssm_backend = ReplaySSMBackend(manager.linear_state_manager.ssm_state.buffer, 16, 3) + ssm_backend = manager.linear_state_manager.ssm_backend + ssm_backend.raw_keys.normal_() + ssm_backend.raw_values.normal_() + ssm_backend.gates.uniform_(-0.1, -0.01) + ssm_backend.betas.uniform_(0.1, 0.9) + ssm_backend.cursors[1] = 2 + active_before = manager.linear_state_manager.ssm_state.buffer[:, 1].clone() + expected = manager.linear_state_manager.ssm_state.buffer[:, 1].float().clone() for i in range(2): - key = cache.raw_keys[:, 1, :, i].float() + key = ssm_backend.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 *= ssm_backend.gates[:, 1, :, i].exp()[..., None, None] + delta = ssm_backend.betas[:, 1, :, i, None] * ( + ssm_backend.raw_values[:, 1, :, i].float() - torch.einsum("lhkv,lhk->lhv", expected, key) ) expected += key[..., None] * delta[..., None, :] expected = expected.to(dtype) - conv_expected = manager.recurrent_state.conv_state.buffer[:, 1, :, 1:4].clone() + conv_expected = manager.linear_state_manager.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) - materialized = cache.snapshot_accepted_state(1) + materialized = ssm_backend.get_accepted_state(1) torch.testing.assert_close(materialized, expected, rtol=state_rtol, atol=2e-5) if save_big: manager.save_big_page_states(reqs, [1], [0]) @@ -611,35 +622,37 @@ def test_checkpoint_restores_pending_history_and_accepted_conv(save_big, monkeyp # a tolerance for FP32 reduction differences crossing BF16 rounding points. torch.testing.assert_close(state_saved, materialized.cpu(), rtol=0, atol=0) torch.testing.assert_close(conv_saved, conv_expected.cpu(), rtol=0, atol=0) - torch.testing.assert_close(manager.recurrent_state.ssm_state.buffer[:, 1], active_before, rtol=0, atol=0) - assert cache.cursors[1].item() == 2 - cache.cursors[2] = 13 # stale history from a previous owner must disappear + torch.testing.assert_close(manager.linear_state_manager.ssm_state.buffer[:, 1], active_before, rtol=0, atol=0) + assert ssm_backend.cursors[1].item() == 2 + ssm_backend.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.recurrent_state.accepted_index[2].item() == 0 - torch.testing.assert_close(manager.recurrent_state.ssm_state.buffer[:, 2], materialized, rtol=0, atol=0) - torch.testing.assert_close(manager.recurrent_state.conv_state.buffer[:, 2, :, :3], conv_expected, rtol=0, atol=0) + assert ssm_backend.cursors[2].item() == 0 + assert manager.linear_state_manager.accepted_state_index[2].item() == 0 + torch.testing.assert_close(manager.linear_state_manager.ssm_state.buffer[:, 2], materialized, rtol=0, atol=0) + torch.testing.assert_close( + manager.linear_state_manager.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, - recurrent_state=manager.recurrent_state, + linear_state_manager=manager.linear_state_manager, ) 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() + ssm_backend.raw_keys[:, 1, :, 0].normal_() + ssm_backend.raw_values[:, 1, :, 0].normal_() + ssm_backend.gates[:, 1, :, 0].uniform_(-0.1, -0.01) + ssm_backend.betas[:, 1, :, 0].uniform_(0.1, 0.9) + ssm_backend.cursors[1] = 1 + key = ssm_backend.raw_keys[:, 1, :, 0].float() key /= (key.square().sum(-1, keepdim=True) + 1e-6).sqrt() expected = active_before.float().clone() - 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 *= ssm_backend.gates[:, 1, :, 0].exp()[..., None, None] + delta = ssm_backend.betas[:, 1, :, 0, None] * ( + ssm_backend.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) @@ -647,26 +660,28 @@ def test_checkpoint_restores_pending_history_and_accepted_conv(save_big, monkeyp ssm_page = torch.empty(helper.ssm_shape, device="cuda", dtype=dtype) helper._write_one_rank(mem, 0, 1, conv_page, ssm_page) torch.testing.assert_close(ssm_page, expected, rtol=state_rtol, atol=2e-5) - torch.testing.assert_close(manager.recurrent_state.ssm_state.buffer[:, 1], active_before, rtol=0, atol=0) - assert cache.cursors[1].item() == 1 - cache.cursors[2] = 7 - manager.recurrent_state.accepted_index[2] = 2 + torch.testing.assert_close(manager.linear_state_manager.ssm_state.buffer[:, 1], active_before, rtol=0, atol=0) + assert ssm_backend.cursors[1].item() == 1 + ssm_backend.cursors[2] = 7 + manager.linear_state_manager.accepted_state_index[2] = 2 helper._read_one_rank(mem, 0, 2, conv_page, ssm_page) - torch.testing.assert_close(manager.recurrent_state.ssm_state.buffer[:, 2], ssm_page, rtol=0, atol=0) - torch.testing.assert_close(manager.recurrent_state.conv_state.buffer[:, 2, :, :3], conv_expected, rtol=0, atol=0) - assert cache.cursors[2].item() == manager.recurrent_state.accepted_index[2].item() == 0 + torch.testing.assert_close(manager.linear_state_manager.ssm_state.buffer[:, 2], ssm_page, rtol=0, atol=0) + torch.testing.assert_close( + manager.linear_state_manager.conv_state.buffer[:, 2, :, :3], conv_expected, rtol=0, atol=0 + ) + assert ssm_backend.cursors[2].item() == manager.linear_state_manager.accepted_state_index[2].item() == 0 manager.init_hybrid_attention_state(SimpleNamespace(req_idx=2)) - assert torch.count_nonzero(manager.recurrent_state.ssm_state.buffer[:, 2]).item() == 0 - assert torch.count_nonzero(manager.recurrent_state.conv_state.buffer[:, 2]).item() == 0 + assert torch.count_nonzero(manager.linear_state_manager.ssm_state.buffer[:, 2]).item() == 0 + assert torch.count_nonzero(manager.linear_state_manager.conv_state.buffer[:, 2]).item() == 0 @pytest.mark.parametrize("dtype", [torch.float32, torch.bfloat16]) @pytest.mark.parametrize("capacity,width", [(4, 1), (4, 3), (8, 4), (16, 5), (32, 4), (64, 4)]) -@pytest.mark.parametrize("projection_mode", ["inline", "precompute"]) -@pytest.mark.parametrize("lower_bound,gate_shift", [(-5.0, 0), (-2.5, -6)]) +@pytest.mark.parametrize("compute_mode", ["inline", "precompute"]) +@pytest.mark.parametrize("gate_lower_bound,gate_shift", [(-5.0, 0), (-2.5, -6)]) @pytest.mark.parametrize("activation_dtype", [torch.float32, torch.bfloat16]) def test_kda_replay_reference_fold_and_graph( - dtype, capacity, width, projection_mode, lower_bound, gate_shift, activation_dtype + dtype, capacity, width, compute_mode, gate_lower_bound, gate_shift, activation_dtype ): test_replay_acceptance_flush_and_graph( width, @@ -674,12 +689,12 @@ def test_kda_replay_reference_fold_and_graph( capacity, dtype, run_config={"BV": 32, "num_warps": 2, "num_stages": 1}, - projection_mode=projection_mode, + compute_mode=compute_mode, gate_shift=gate_shift, activation_dtype=activation_dtype, head_geometry=(2, 2), kda=True, - lower_bound=lower_bound, + gate_lower_bound=gate_lower_bound, ) @@ -700,7 +715,7 @@ def test_kda_tuning_candidates_match_reference(config, mode): 4, torch.bfloat16, run_config=config, - projection_mode=mode, + compute_mode=mode, head_geometry=(2, 2), kda=True, gate_shift=-3, diff --git a/unit_tests/common/basemodel/triton_kernel/linear_att/test_replayssm_execution.py b/unit_tests/common/basemodel/triton_kernel/linear_att/test_replayssm_execution.py index e2ba53add1..0dac93dad9 100644 --- a/unit_tests/common/basemodel/triton_kernel/linear_att/test_replayssm_execution.py +++ b/unit_tests/common/basemodel/triton_kernel/linear_att/test_replayssm_execution.py @@ -1,11 +1,11 @@ import pytest import torch -from lightllm.common.basemodel.triton_kernel.linear_att.replayssm import ReplaySSMCache +from lightllm.common.basemodel.triton_kernel.linear_att.replayssm import ReplaySSMBackend from lightllm.common.basemodel.triton_kernel.linear_att.mtp_state_params import ( build_dynamic_mtp_linear_att_state_params, ) -from lightllm.common.recurrent_state import RecurrentState +from lightllm.common.linear_state import LinearStateManager from lightllm.common.req_manager.linear_att import ReqManagerForMamba pytestmark = pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") @@ -19,7 +19,7 @@ def test_request_manager_accept_with_replay_graph(dtype, capture_metadata, varle torch.manual_seed(827) config = {"BV": 32, "num_warps": 1, "num_stages": 1} initial = torch.randn(3, 5, 4, 128, 128, device="cuda", dtype=dtype) * 0.01 - caches = [ + ssm_backends = [ cls( initial.clone(), 8, @@ -27,15 +27,15 @@ def test_request_manager_accept_with_replay_graph(dtype, capture_metadata, varle torch.bfloat16, num_key_heads=2, kda=kda, - projection_mode="precompute", + compute_mode="precompute", run_config=config, ) - for cls in [ReplaySSMCache, ReplaySSMCache] + for cls in [ReplaySSMBackend, ReplaySSMBackend] ] manager = object.__new__(ReqManagerForMamba) - manager.recurrent_state = object.__new__(RecurrentState) - manager.recurrent_state.updates = caches[1] - manager.recurrent_state.accepted_index = torch.zeros(5, device="cuda", dtype=torch.int32) + manager.linear_state_manager = object.__new__(LinearStateManager) + manager.linear_state_manager.ssm_backend = ssm_backends[1] + manager.linear_state_manager.accepted_state_index = torch.zeros(5, device="cuda", dtype=torch.int32) rows = torch.full((12,), 4, device="cuda", dtype=torch.int32) mtp = torch.zeros_like(rows) flags = torch.zeros_like(rows) @@ -72,53 +72,53 @@ def inputs(iteration): accepted.copy_(accepted.new_tensor(table)) return ids, lengths, counts, logical - def metadata(cache): - cu, reqs, _ = build_dynamic_mtp_linear_att_state_params(rows, mtp, manager.recurrent_state.accepted_index, 4) - return cu, reqs, cache.prepare_decode(reqs, cu) + def metadata(ssm_backend): + cu, reqs = build_dynamic_mtp_linear_att_state_params(rows, mtp, 4) + return cu, reqs, ssm_backend.prepare_decode(reqs, cu) def forward_accept(index, cu, reqs): - cache = caches[index] - out = [cache.forward(layer, q, k, v, a, beta, alog, bias, reqs, cu) for layer in range(3)] + ssm_backend = ssm_backends[index] + out = [ssm_backend.forward(layer, q, k, v, a, beta, alog, bias, reqs, cu)[0] for layer in range(3)] if index == 0: - cache.accept_updates(reqs, accepted) + ssm_backend.accept_updates(reqs, accepted) else: manager.update_mtp_state(starts, rows, mtp, flags, 4) return out inputs(0) graphs, captured, outputs = [], [], [] - for index, cache in enumerate(caches): - cu, reqs, pos = metadata(cache) + for index, ssm_backend in enumerate(ssm_backends): + cu, reqs, pos = metadata(ssm_backend) forward_accept(index, cu, reqs) - cache.state.copy_(initial) - cache.cursors.zero_() - manager.recurrent_state.accepted_index.zero_() + ssm_backend.state.copy_(initial) + ssm_backend.cursors.zero_() + manager.linear_state_manager.accepted_state_index.zero_() g = torch.cuda.CUDAGraph() with torch.cuda.graph(g): if capture_metadata: - cu, reqs, pos = metadata(cache) + cu, reqs, pos = metadata(ssm_backend) out = forward_accept(index, cu, reqs) graphs.append(g) captured.append((cu, reqs, pos)) outputs.append(out) - for cache in caches: - cache.state.copy_(initial) - cache.cursors.zero_() - manager.recurrent_state.accepted_index.zero_() + for ssm_backend in ssm_backends: + ssm_backend.state.copy_(initial) + ssm_backend.cursors.zero_() + manager.linear_state_manager.accepted_state_index.zero_() folds = 0 for iteration in range(32): ids, lengths, counts, logical = inputs(iteration) for tensor in [q, k, v, beta]: tensor.normal_() - for index, cache in enumerate(caches): + for index, ssm_backend in enumerate(ssm_backends): if not capture_metadata: - for dst, src in zip(captured[index], metadata(cache)): + for dst, src in zip(captured[index], metadata(ssm_backend)): dst.copy_(src) graphs[index].replay() torch.cuda.synchronize() for expected, actual in zip(captured[0], captured[1]): assert torch.equal(actual, expected) - assert torch.equal(caches[0].cursors, caches[1].cursors) + assert torch.equal(ssm_backends[0].cursors, ssm_backends[1].cursors) positions = captured[0][2].cpu().tolist() folds += sum(positions[j] % 16 + length > 8 for j, length in enumerate(lengths)) for expected, actual in zip(outputs[0], outputs[1]): @@ -126,13 +126,13 @@ def forward_accept(index, cu, reqs): actual[:, :logical].float(), expected[:, :logical].float(), atol=0.006, rtol=0.02 ) torch.testing.assert_close( - caches[1].state.float(), caches[0].state.float(), atol=0.002, rtol=torch.finfo(dtype).eps + ssm_backends[1].state.float(), ssm_backends[0].state.float(), atol=0.002, rtol=torch.finfo(dtype).eps ) - table = manager.recurrent_state.accepted_index.cpu().tolist() + table = manager.linear_state_manager.accepted_state_index.cpu().tolist() for req, count in zip(ids, counts): assert table[req] == count - 1 assert table[2] == table[4] == 0 - assert caches[1].cursors[4].item() == 0 + assert ssm_backends[1].cursors[4].item() == 0 assert folds > 0 @@ -143,8 +143,8 @@ def forward_accept(index, cu, reqs): def test_decode_accept_order_preserves_forward_snapshot(dtype, capture_metadata, stagger, kda): torch.manual_seed(731) initial = torch.randn(3, 3, 4, 128, 128, device="cuda", dtype=dtype) * 0.01 - caches = [ - ReplaySSMCache( + ssm_backends = [ + ReplaySSMBackend( initial.clone(), 8, 1, @@ -155,8 +155,8 @@ def test_decode_accept_order_preserves_forward_snapshot(dtype, capture_metadata, ) for _ in range(2) ] - for cache in caches: - cache.fold_programs = 3 + for ssm_backend in ssm_backends: + ssm_backend.fold_programs = 3 reqs = torch.tensor([0, 1, 2], device="cuda", dtype=torch.int32) q = torch.randn(3, 2, 128, device="cuda", dtype=dtype) k = torch.randn_like(q) @@ -166,35 +166,35 @@ def test_decode_accept_order_preserves_forward_snapshot(dtype, capture_metadata, alog = torch.zeros(4, device="cuda") bias = torch.zeros(4 * 128 if kda else 4, device="cuda") - def forward(cache): - return [cache.forward(layer, q, k, v, a, beta, alog, bias, reqs) for layer in range(3)] + def forward(ssm_backend): + return [ssm_backend.forward(layer, q, k, v, a, beta, alog, bias, reqs)[0] for layer in range(3)] - def warm(cache): - cache.prepare_decode(reqs) - out = forward(cache) - cache.accept_updates(reqs) + def warm(ssm_backend): + ssm_backend.prepare_decode(reqs) + out = forward(ssm_backend) + ssm_backend.accept_updates(reqs) return out graphs, positions, outputs = [], [], [] - for early, cache in enumerate(caches): - warm(cache) - cache.state.copy_(initial) - cache.cursors.zero_() + for early, ssm_backend in enumerate(ssm_backends): + warm(ssm_backend) + ssm_backend.state.copy_(initial) + ssm_backend.cursors.zero_() if stagger: reqs.copy_(reqs.new_tensor([0, 2, 2])) for _ in range(3): - warm(cache) + warm(ssm_backend) reqs.copy_(reqs.new_tensor([0, 1, 2])) pos = torch.zeros_like(reqs) graph = torch.cuda.CUDAGraph() with torch.cuda.graph(graph): if capture_metadata: - pos = cache.prepare_decode(reqs) + pos = ssm_backend.prepare_decode(reqs) if early: - cache.accept_updates(reqs) - out = forward(cache) + ssm_backend.accept_updates(reqs) + out = forward(ssm_backend) if capture_metadata and not early: - cache.accept_updates(reqs) + ssm_backend.accept_updates(reqs) graphs.append(graph) positions.append(pos) outputs.append(out) @@ -203,23 +203,23 @@ def warm(cache): reqs.copy_(reqs.new_tensor([0, 1, 2] if iteration % 2 == 0 else [1, 0, 2])) for tensor in [q, k, v, beta]: tensor.normal_() - old = caches[0].cursors[reqs.long()].clone() - for early, (cache, graph, pos) in enumerate(zip(caches, graphs, positions)): + old = ssm_backends[0].cursors[reqs.long()].clone() + for early, (ssm_backend, graph, pos) in enumerate(zip(ssm_backends, graphs, positions)): if not capture_metadata: - pos.copy_(cache.prepare_decode(reqs)) + pos.copy_(ssm_backend.prepare_decode(reqs)) if early: - cache.accept_updates(reqs) + ssm_backend.accept_updates(reqs) graph.replay() if not capture_metadata and not early: - cache.accept_updates(reqs) + ssm_backend.accept_updates(reqs) torch.cuda.synchronize() assert torch.equal(positions[0], old) assert torch.equal(positions[0], positions[1]) - assert torch.equal(caches[0].cursors, caches[1].cursors) + assert torch.equal(ssm_backends[0].cursors, ssm_backends[1].cursors) for expected, actual in zip(outputs[0], outputs[1]): torch.testing.assert_close(actual.float(), expected.float(), atol=0.006, rtol=0.02) torch.testing.assert_close( - caches[1].state.float(), caches[0].state.float(), atol=0.002, rtol=torch.finfo(dtype).eps + ssm_backends[1].state.float(), ssm_backends[0].state.float(), atol=0.002, rtol=torch.finfo(dtype).eps ) folded = any(n % 16 + 1 > 8 for n in old[:2].cpu().tolist()) saw_fold |= folded @@ -238,8 +238,8 @@ def test_cross_layer_fold_with_state_offset_above_int32(): sparse = torch.empty_strided(shape, strides, device="cuda", dtype=dense.dtype) sparse.copy_(dense) config = {"BV": 32, "num_warps": 1, "num_stages": 1} - caches = [ - ReplaySSMCache(state, 8, 4, torch.bfloat16, num_key_heads=4, run_config=config) for state in [dense, sparse] + ssm_backends = [ + ReplaySSMBackend(state, 8, 4, torch.bfloat16, num_key_heads=4, run_config=config) for state in [dense, sparse] ] q = torch.randn(1, 8, 4, 128, device="cuda", dtype=torch.bfloat16) k = torch.randn_like(q) @@ -251,22 +251,22 @@ def test_cross_layer_fold_with_state_offset_above_int32(): reqs = torch.tensor([0, 2], device="cuda", dtype=torch.int32) cu = reqs.new_tensor([0, 4, 8]) - def step(cache, accepted): - pos = cache.prepare_decode(reqs, cu) - out = [cache.forward(layer, q, k, v, a, beta, alog, bias, reqs, cu) for layer in range(48)] - cache.accept_updates(reqs, accepted) + def step(ssm_backend, accepted): + pos = ssm_backend.prepare_decode(reqs, cu) + out = [ssm_backend.forward(layer, q, k, v, a, beta, alog, bias, reqs, cu)[0] for layer in range(48)] + ssm_backend.accept_updates(reqs, accepted) return pos, out for count in [2, 3]: accepted = reqs.new_tensor([count - 1, -1, -1]) - for cache in caches: - step(cache, accepted) + for ssm_backend in ssm_backends: + step(ssm_backend, accepted) reqs.copy_(reqs.new_tensor([0, 1])) - expected_pos, expected_out = step(caches[0], reqs.new_tensor([0, 0, -1])) - actual_pos, actual_out = step(caches[1], reqs.new_tensor([0, 0, -1])) + expected_pos, expected_out = step(ssm_backends[0], reqs.new_tensor([0, 0, -1])) + actual_pos, actual_out = step(ssm_backends[1], reqs.new_tensor([0, 0, -1])) torch.cuda.synchronize() assert torch.equal(actual_pos, expected_pos) - assert torch.equal(caches[0].cursors, caches[1].cursors) + assert torch.equal(ssm_backends[0].cursors, ssm_backends[1].cursors) for layer in range(48): torch.testing.assert_close(actual_out[layer].float(), expected_out[layer].float(), atol=0.006, rtol=0.02) torch.testing.assert_close( @@ -281,8 +281,8 @@ def test_small_batch_matches_larger_graph_padding(batch, dtype, kda): torch.manual_seed(1087 + batch) initial = torch.randn(2, 14, 4, 128, 128, device="cuda", dtype=dtype) * 0.01 config = {"BV": 32, "num_warps": 1, "num_stages": 1} - caches = [ - ReplaySSMCache(initial.clone(), 8, 4, num_key_heads=2, kda=kda, projection_mode="precompute", run_config=config) + ssm_backends = [ + ReplaySSMBackend(initial.clone(), 8, 4, num_key_heads=2, kda=kda, compute_mode="precompute", run_config=config) for _ in range(2) ] reqs = torch.arange(batch, device="cuda", dtype=torch.int32) @@ -298,12 +298,12 @@ def test_small_batch_matches_larger_graph_padding(batch, dtype, kda): bias = torch.zeros(4 * 128 if kda else 4, device="cuda") accepted = torch.zeros(14, device="cuda", dtype=torch.int32) graphs, outputs = [], [] - for cache, ids, offsets in zip(caches, [reqs, padded], [cu, padded_cu]): + for ssm_backend, ids, offsets in zip(ssm_backends, [reqs, padded], [cu, padded_cu]): def step(): - cache.prepare_decode(ids, offsets) - out = [cache.forward(layer, q, k, v, a, beta, alog, bias, ids, offsets) for layer in range(2)] - cache.accept_updates(ids, accepted) + ssm_backend.prepare_decode(ids, offsets) + out = [ssm_backend.forward(layer, q, k, v, a, beta, alog, bias, ids, offsets)[0] for layer in range(2)] + ssm_backend.accept_updates(ids, accepted) return out step() @@ -312,8 +312,8 @@ def step(): out = step() graphs.append(graph) outputs.append(out) - cache.state.copy_(initial) - cache.cursors.zero_() + ssm_backend.state.copy_(initial) + ssm_backend.cursors.zero_() for iteration in range(32): accepted[:batch] = (torch.arange(batch, device="cuda") + iteration) % 4 for value in [q, k, v, beta]: @@ -322,7 +322,7 @@ def step(): graph.replay() for small, padded_out in zip(outputs[0], outputs[1]): torch.testing.assert_close(small.float(), padded_out.float(), atol=0.006, rtol=0.02) - assert torch.equal(caches[0].cursors, caches[1].cursors) + assert torch.equal(ssm_backends[0].cursors, ssm_backends[1].cursors) torch.testing.assert_close( - caches[0].state.float(), caches[1].state.float(), atol=0.002, rtol=torch.finfo(dtype).eps + ssm_backends[0].state.float(), ssm_backends[1].state.float(), atol=0.002, rtol=torch.finfo(dtype).eps ) diff --git a/unit_tests/common/basemodel/triton_kernel/linear_att/test_replayssm_grouped_keys.py b/unit_tests/common/basemodel/triton_kernel/linear_att/test_replayssm_grouped_keys.py index ea4a2795a7..7546fd38f9 100644 --- a/unit_tests/common/basemodel/triton_kernel/linear_att/test_replayssm_grouped_keys.py +++ b/unit_tests/common/basemodel/triton_kernel/linear_att/test_replayssm_grouped_keys.py @@ -3,8 +3,8 @@ import pytest import torch -from lightllm.common.basemodel.triton_kernel.linear_att.replayssm import ReplaySSMCache -from lightllm.common.basemodel.triton_kernel.linear_att.replayssm_compact import CompactSSMCache +from lightllm.common.basemodel.triton_kernel.linear_att.replayssm import ReplaySSMBackend +from lightllm.common.basemodel.triton_kernel.linear_att.compact_ssm import CompactSSMBackend pytestmark = pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") @@ -18,13 +18,13 @@ def test_grouped_keys_graph_fold_partial_accept(mode, group, dtype): hv = h * group state = torch.randn(layers, slots, hv, kd, vd, device="cuda", dtype=dtype) * 0.01 - def make_cache(key_heads): + def make_ssm_backend(key_heads): if mode == "replay": - return ReplaySSMCache(state.clone(), 16, width, num_key_heads=key_heads) - return CompactSSMCache(state.clone(), width, torch.bfloat16, num_key_heads=key_heads) + return ReplaySSMBackend(state.clone(), 16, width, num_key_heads=key_heads) + return CompactSSMBackend(state.clone(), width, torch.bfloat16, num_key_heads=key_heads) - caches = [make_cache(hv), make_cache(h)] - assert caches[0].keys.numel() == caches[1].keys.numel() * group + ssm_backends = [make_ssm_backend(hv), make_ssm_backend(h)] + assert ssm_backends[0].keys.numel() == ssm_backends[1].keys.numel() * group # Packed strides, out-of-order requests, an empty real request, nonempty HOLD. packed = torch.randn(1, 8, (2 * h + hv) * kd, device="cuda", dtype=torch.bfloat16) q, k, v = [x.view(1, 8, heads, kd) for x, heads in zip(packed.split([h * kd, h * kd, hv * kd], -1), [h, h, hv])] @@ -36,23 +36,23 @@ def make_cache(key_heads): cu = torch.tensor([0, 4, 7, 7, 8], device="cuda", dtype=torch.int32) accepted = torch.tensor([0, -1, 0, -1, -1], device="cuda", dtype=torch.int32) - def step(cache): - cache.prepare_decode(reqs, cu) - out = [cache.forward(layer, q, k, v, a, b, alog, bias, reqs, cu) for layer in range(layers)] - cache.accept_updates(reqs, accepted) + def step(ssm_backend): + ssm_backend.prepare_decode(reqs, cu) + out = [ssm_backend.forward(layer, q, k, v, a, b, alog, bias, reqs, cu)[0] for layer in range(layers)] + ssm_backend.accept_updates(reqs, accepted) return out graphs, outputs = [], [] - for cache in caches: - step(cache) - cache.state.copy_(state) + for ssm_backend in ssm_backends: + step(ssm_backend) + ssm_backend.state.copy_(state) if mode == "replay": - cache.cursors.zero_() + ssm_backend.cursors.zero_() reqs.fill_(4) torch.cuda.synchronize() graph = torch.cuda.CUDAGraph() with torch.cuda.graph(graph): - outputs.append(step(cache)) + outputs.append(step(ssm_backend)) graphs.append(graph) reqs.copy_(torch.tensor([2, 0, 1, 4], device="cuda", dtype=torch.int32)) @@ -68,14 +68,14 @@ def step(cache): torch.testing.assert_close(left, right, atol=0, rtol=0) assert torch.count_nonzero(right[:, -1]) == 0 if iteration in (18, 39, 66): - for cache in caches: - cache.merge_accepted_updates(reqs) - torch.testing.assert_close(caches[0].state, caches[1].state, atol=0, rtol=0) - torch.testing.assert_close(caches[1].state[:, [1, 3, 4]], state[:, [1, 3, 4]], atol=0, rtol=0) + for ssm_backend in ssm_backends: + ssm_backend.merge_accepted_updates(reqs) + torch.testing.assert_close(ssm_backends[0].state, ssm_backends[1].state, atol=0, rtol=0) + torch.testing.assert_close(ssm_backends[1].state[:, [1, 3, 4]], state[:, [1, 3, 4]], atol=0, rtol=0) if mode == "replay": - torch.testing.assert_close(caches[0].cursors, caches[1].cursors, atol=0, rtol=0) + torch.testing.assert_close(ssm_backends[0].cursors, ssm_backends[1].cursors, atol=0, rtol=0) if iteration == 40: # Request-slot reuse / checkpoint restore must not resurrect history. - for cache in caches: - cache.clear_history(2) - cache.state[:, 2].copy_(state[:, 2]) + for ssm_backend in ssm_backends: + ssm_backend.clear_history(2) + ssm_backend.state[:, 2].copy_(state[:, 2]) diff --git a/unit_tests/common/basemodel/triton_kernel/linear_att/test_replayssm_parallel.py b/unit_tests/common/basemodel/triton_kernel/linear_att/test_replayssm_parallel.py index 0a2248080c..7d869ac2c4 100644 --- a/unit_tests/common/basemodel/triton_kernel/linear_att/test_replayssm_parallel.py +++ b/unit_tests/common/basemodel/triton_kernel/linear_att/test_replayssm_parallel.py @@ -1,7 +1,7 @@ import pytest import torch -from lightllm.common.basemodel.triton_kernel.linear_att.replayssm import ReplaySSMCache +from lightllm.common.basemodel.triton_kernel.linear_att.replayssm import ReplaySSMBackend pytestmark = pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") @@ -11,48 +11,48 @@ def test_derived_fold_commits_only_accepted_history(capacity): torch.manual_seed(731) state = torch.randn(2, 5, 4, 128, 128, device="cuda", dtype=torch.bfloat16) * 0.02 initial = state.clone() - cache = ReplaySSMCache(state, capacity, 5, num_key_heads=2, projection_mode="parallel", fold_method="derived") - cache.keys.fill_(float("nan")) - cache.deltas.fill_(float("nan")) - cache.gates.fill_(float("nan")) - cache.raw_keys.fill_(float("nan")) - cache.raw_values.fill_(float("nan")) - cache.betas.fill_(float("nan")) + ssm_backend = ReplaySSMBackend(state, capacity, 5, num_key_heads=2, compute_mode="parallel", fold_method="derived") + ssm_backend.keys.fill_(float("nan")) + ssm_backend.deltas.fill_(float("nan")) + ssm_backend.gates.fill_(float("nan")) + ssm_backend.raw_keys.fill_(float("nan")) + ssm_backend.raw_values.fill_(float("nan")) + ssm_backend.betas.fill_(float("nan")) # Fold request 2 from the second ring half; request 0 has an empty span. count = capacity - 2 - cache.cursors[2] = 2 * capacity + count - cache.cursors[0] = count - cache.keys[:, 2, :, capacity : capacity + count].normal_(0, 0.1) - cache.deltas[:, 2, :, capacity : capacity + count].normal_(0, 0.1) - cache.gates[:, 2, :, capacity : capacity + count].fill_(-0.1) + ssm_backend.cursors[2] = 2 * capacity + count + ssm_backend.cursors[0] = count + ssm_backend.keys[:, 2, :, capacity : capacity + count].normal_(0, 0.1) + ssm_backend.deltas[:, 2, :, capacity : capacity + count].normal_(0, 0.1) + ssm_backend.gates[:, 2, :, capacity : capacity + count].fill_(-0.1) expected = initial[:, 2].float() for j in range(count): - keys = cache.keys[:, 2, :, capacity + j].repeat_interleave(2, dim=1) - delta = cache.deltas[:, 2, :, capacity + j] - decay = cache.gates[:, 2, :, capacity + j].exp() + keys = ssm_backend.keys[:, 2, :, capacity + j].repeat_interleave(2, dim=1) + delta = ssm_backend.deltas[:, 2, :, capacity + j] + decay = ssm_backend.gates[:, 2, :, capacity + j].exp() expected = expected * decay[:, :, None, None] + keys[:, :, :, None] * delta[:, :, None, :] reqs = torch.tensor([2, 0, 4], dtype=torch.int32, device="cuda") cu = torch.tensor([0, 5, 5, 10], dtype=torch.int32, device="cuda") - cache.prepare_decode(reqs, cu) + ssm_backend.prepare_decode(reqs, cu) assert torch.isfinite(state).all() torch.testing.assert_close(state[:, 2], expected.bfloat16(), atol=2e-4, rtol=0.008) torch.testing.assert_close(state[:, [0, 1, 3, 4]], initial[:, [0, 1, 3, 4]], atol=0, rtol=0) - assert cache.cursors[2] == 0 - assert cache.cursors[0] == count - assert cache.cursors[4] == 0 + assert ssm_backend.cursors[2] == 0 + assert ssm_backend.cursors[0] == count + assert ssm_backend.cursors[4] == 0 @pytest.mark.parametrize("fold", ["raw", "derived"]) def test_parallel_rejects_unsupported_shapes(fold): state = torch.zeros(1, 3, 4, 32, 64, device="cuda", dtype=torch.bfloat16) with pytest.raises(AssertionError, match="128x128"): - ReplaySSMCache(state, 16, 5, num_key_heads=2, projection_mode="parallel", fold_method=fold) + ReplaySSMBackend(state, 16, 5, num_key_heads=2, compute_mode="parallel", fold_method=fold) def test_derived_fold_requires_parallel_solver(): state = torch.zeros(1, 3, 4, 128, 128, device="cuda", dtype=torch.bfloat16) with pytest.raises(AssertionError): - ReplaySSMCache(state, 16, 5, num_key_heads=2, fold_method="derived") + ReplaySSMBackend(state, 16, 5, num_key_heads=2, fold_method="derived") @pytest.mark.parametrize("width", [2, 3, 5, 8, 16]) @@ -66,7 +66,7 @@ def test_parallel_raw_against_sequential_reference(width): dims=(128, 128), capacity=16, dtype=torch.bfloat16, - projection_mode="parallel", + compute_mode="parallel", run_config={"BV": 128, "num_warps": 2, "num_stages": 1}, ) @@ -79,8 +79,17 @@ def test_parallel_derived_continuous_state(capacity, width, capture): torch.manual_seed(731) state = torch.randn(2, 4, 4, 128, 128, device="cuda", dtype=torch.bfloat16) * 0.02 initial = state.clone() - cache = ReplaySSMCache(state, capacity, width, num_key_heads=2, projection_mode="parallel", fold_method="derived") - for history in (cache.keys, cache.deltas, cache.gates, cache.raw_keys, cache.raw_values, cache.betas): + ssm_backend = ReplaySSMBackend( + state, capacity, width, num_key_heads=2, compute_mode="parallel", fold_method="derived" + ) + for history in ( + ssm_backend.keys, + ssm_backend.deltas, + ssm_backend.gates, + ssm_backend.raw_keys, + ssm_backend.raw_values, + ssm_backend.betas, + ): history.fill_(float("nan")) reqs = torch.tensor([2, 0, 3], device="cuda", dtype=torch.int32) lengths = [width, width - 1] @@ -95,17 +104,17 @@ def test_parallel_derived_continuous_state(capacity, width, capture): accepted = torch.zeros(4, device="cuda", dtype=torch.int32) def step(): - cache.prepare_decode(reqs, cu) - outputs = [cache.forward(layer, q, k, v, a, b, alog, bias, reqs, cu) for layer in range(2)] - cache.accept_updates(reqs, accepted) + ssm_backend.prepare_decode(reqs, cu) + outputs = [ssm_backend.forward(layer, q, k, v, a, b, alog, bias, reqs, cu)[0] for layer in range(2)] + ssm_backend.accept_updates(reqs, accepted) return outputs # Compile before capture, then reset the live state and capture only HOLD rows. step() state.copy_(initial) - cache.cursors.zero_() + ssm_backend.cursors.zero_() if capture: - reqs.fill_(cache.hold) + reqs.fill_(ssm_backend.hold_req_id) torch.cuda.synchronize() graph = torch.cuda.CUDAGraph() with torch.cuda.graph(graph): @@ -145,19 +154,19 @@ def step(): assert torch.count_nonzero(outputs[layer][0, -1]) == 0 if iteration == capacity: - saved_state, saved_cursors = state.clone(), cache.cursors.clone() - snapshot = cache.snapshot_accepted_state(2) + saved_state, saved_cursors = state.clone(), ssm_backend.cursors.clone() + snapshot = ssm_backend.get_accepted_state(2) torch.testing.assert_close(snapshot.float(), expected[:, 2], atol=0.002, rtol=0.02) torch.testing.assert_close(state, saved_state, atol=0, rtol=0) - torch.testing.assert_close(cache.cursors, saved_cursors, atol=0, rtol=0) + torch.testing.assert_close(ssm_backend.cursors, saved_cursors, atol=0, rtol=0) # Restore the canonical checkpoint while another request keeps its history. - cache.clear_history(2) + ssm_backend.clear_history(2) state[:, 2].copy_(snapshot) expected[:, 2].copy_(snapshot.float()) pending[0] = 0 assert min(folds) >= 2 # Exercise both ring halves, including after restore. - cache.merge_accepted_updates(reqs[:2]) + ssm_backend.merge_accepted_updates(reqs[:2]) torch.testing.assert_close(state.float(), expected.bfloat16().float(), atol=0.002, rtol=0.02) torch.testing.assert_close(state[:, [1, 3]], initial[:, [1, 3]], atol=0, rtol=0) - assert torch.count_nonzero(cache.cursors) == 0 + assert torch.count_nonzero(ssm_backend.cursors) == 0 diff --git a/unit_tests/common/basemodel/triton_kernel/linear_att/test_ssm_autotune.py b/unit_tests/common/basemodel/triton_kernel/linear_att/test_ssm_autotune.py index f88689f3df..e824ad77c1 100644 --- a/unit_tests/common/basemodel/triton_kernel/linear_att/test_ssm_autotune.py +++ b/unit_tests/common/basemodel/triton_kernel/linear_att/test_ssm_autotune.py @@ -1,27 +1,27 @@ import collections import math +from dataclasses import asdict import pytest import torch from lightllm.common.basemodel.triton_kernel.linear_att import ssm_autotune as tuning -from lightllm.common.basemodel.triton_kernel.linear_att.replayssm import ReplaySSMCache -from lightllm.common.basemodel.triton_kernel.linear_att.replayssm_compact import CompactSSMCache -from lightllm.common.recurrent_state import RecurrentState -from lightllm.common.recurrent_state_types import SSMRunConfig +from lightllm.common.basemodel.triton_kernel.linear_att.replayssm import ReplaySSMBackend +from lightllm.common.basemodel.triton_kernel.linear_att.compact_ssm import CompactSSMBackend +from lightllm.common.linear_state import LinearStateManager from lightllm.common.kernel_config import KernelConfigs from lightllm.common import kernel_config as kernel_config_module from lightllm.common.triton_utils import autotuner as autotuner_module from lightllm.common.triton_utils.autotuner import Autotuner, AutotuneKernelType -def make_inputs(mode, dtype=torch.bfloat16, device="cpu", width=4, projection_mode="inline", kda=False): +def make_inputs(mode, dtype=torch.bfloat16, device="cpu", width=4, compute_mode="inline", kda=False): layers, slots, batch, h, hv, kd, vd = 2, 9, 3, 2, 4, 32, 64 state = torch.randn(layers, slots, hv, kd, vd, dtype=dtype, device=device) * 0.01 if mode == "replay": - cache = ReplaySSMCache(state, 8, width, num_key_heads=h, projection_mode=projection_mode, kda=kda) + ssm_backend = ReplaySSMBackend(state, 8, width, num_key_heads=h, compute_mode=compute_mode, kda=kda) else: - cache = CompactSSMCache(state, width, torch.bfloat16, num_key_heads=h, kda=kda) + ssm_backend = CompactSSMBackend(state, width, torch.bfloat16, num_key_heads=h, kda=kda) tokens = batch * width packed = torch.randn(1, tokens, 2 * h * kd + hv * vd, dtype=torch.bfloat16, device=device) q, k, v = [ @@ -31,7 +31,7 @@ def make_inputs(mode, dtype=torch.bfloat16, device="cpu", width=4, projection_mo gates = torch.randn(tokens, hv * (kd + 1) if kda else hv * 2, device=device, dtype=q.dtype) a, b = gates[:, :-hv], gates[:, -hv:] return dict( - cache=cache, + ssm_backend=ssm_backend, mode=mode, q=q, k=k, @@ -44,23 +44,23 @@ def make_inputs(mode, dtype=torch.bfloat16, device="cpu", width=4, projection_mo ) -def tensors(cache): - return {name: x for name, x in vars(cache).items() if isinstance(x, torch.Tensor)} +def tensors(ssm_backend): + return {name: x for name, x in vars(ssm_backend).items() if isinstance(x, torch.Tensor)} def test_replay_projection_has_separate_autotune_cache(): inline = make_inputs("replay") - precompute = make_inputs("replay", projection_mode="precompute") + precompute = make_inputs("replay", compute_mode="precompute") assert tuning.select_config._static_key(**inline) != tuning.select_config._static_key(**precompute) _, rebuilt = tuning.rebuild_inputs(**precompute) - assert rebuilt["cache"].projection_mode == "precompute" - assert rebuilt["cache"].keys.dtype == precompute["cache"].keys.dtype == torch.bfloat16 + assert rebuilt["ssm_backend"].compute_mode == "precompute" + assert rebuilt["ssm_backend"].keys.dtype == precompute["ssm_backend"].keys.dtype == torch.bfloat16 assert tuning.select_config._static_key(**precompute) == tuning.select_config._static_key(**rebuilt) key = tuning.select_config._static_key(**precompute) legacy = dict(key, v=2) assert KernelConfigs.get_config_file_name(key) != KernelConfigs.get_config_file_name(legacy) assert len(KernelConfigs.get_config_file_name(key).encode()) <= 255 - precompute["cache"].keys = precompute["cache"].keys.float() + precompute["ssm_backend"].keys = precompute["ssm_backend"].keys.float() assert tuning.select_config._static_key(**precompute) != key @@ -75,20 +75,21 @@ def isolate_tuner(monkeypatch, tmp_path): monkeypatch.setattr(tuner, "cached_configs", {}) monkeypatch.setattr(tuner, "fast_match_configs", collections.defaultdict(dict)) monkeypatch.setattr(tuner, "warmuped_configs_set", set()) + monkeypatch.setattr(tuner, "_oom_keys", set()) -@pytest.mark.parametrize("mode", ["gdn", "replay"]) +@pytest.mark.parametrize("mode", ["compact", "replay"]) @pytest.mark.parametrize("dtype", [torch.float32, torch.bfloat16]) @pytest.mark.parametrize("kda", [False, True], ids=["gdn", "kda"]) def test_rebuild_has_real_requests_and_independent_packed_inputs(mode, dtype, kda): inputs = make_inputs(mode, dtype, kda=kda) - cache = inputs["cache"] - for x in tensors(cache).values(): + ssm_backend = inputs["ssm_backend"] + for x in tensors(ssm_backend).values(): x.fill_(1) args, rebuilt = tuning.rebuild_inputs(**inputs) assert not args - scratch = rebuilt["cache"] - assert scratch.state.shape[:2] == (2, 4) # Active batch + HOLD, not the serving pool. + scratch = rebuilt["ssm_backend"] + assert scratch.state.shape[:2] == (1, 4) # One layer, active batch + HOLD. assert scratch._config_is_fixed assert rebuilt["workload"][0].tolist() == [0, 1, 2] assert rebuilt["workload"][1].tolist() == [3, 3, 3, 0] @@ -100,7 +101,7 @@ def test_rebuild_has_real_requests_and_independent_packed_inputs(mode, dtype, kd assert tuning.select_config._run_key(**inputs) == tuning.select_config._run_key(**rebuilt) for x in tensors(scratch).values(): x.zero_() - for x in tensors(cache).values(): + for x in tensors(ssm_backend).values(): assert torch.all(x == 1) key = tuning.select_config._static_key(**inputs) assert len(KernelConfigs.get_config_file_name(key).encode()) <= 255 @@ -114,34 +115,34 @@ def test_rebuild_has_real_requests_and_independent_packed_inputs(mode, dtype, kd @pytest.mark.parametrize("level", [0, 1, 2, 3]) def test_layout_is_fixed_before_graph_capture_and_cached_configs_do_not_execute(monkeypatch, level): - inputs = make_inputs("gdn") - cache = inputs["cache"] + inputs = make_inputs("compact") + ssm_backend = inputs["ssm_backend"] config = {"BV": 16, "num_warps": 1, "num_stages": 2} tuner = tuning.select_config key = autotuner_module.frozendict(tuner._static_key(**inputs)) tuner.cached_configs[key] = {str(tuner._run_key(**inputs)): config} monkeypatch.setattr(autotuner_module, "get_triton_autotune_level", lambda: level) monkeypatch.setattr(tuner, "_bench", lambda *args, **kwargs: pytest.fail("cached layout must not benchmark")) - default = cache.run_config.copy() + default = ssm_backend.run_config.copy() with Autotuner.autotune_warmup(AutotuneKernelType.GENERAL): - tuning.configure_cache(**inputs) - assert cache.run_config == default and not cache._config_is_fixed + tuning.resolve_run_config(**inputs) + assert ssm_backend.run_config == default and not ssm_backend._config_is_fixed with Autotuner.autotune_warmup(AutotuneKernelType.DECODE_ATTENTION): - selected = tuning.configure_cache(**inputs) + selected = tuning.resolve_run_config(**inputs) assert selected == (default if level == 3 else config) - assert cache.run_config == default and not cache._config_is_fixed + assert ssm_backend.run_config == default and not ssm_backend._config_is_fixed # Reusing the same bucket must not tune again. monkeypatch.setattr(tuning, "select_config", lambda *args, **kwargs: pytest.fail("layout changed after capture")) with Autotuner.autotune_warmup(AutotuneKernelType.DECODE_ATTENTION): - tuning.configure_cache(**inputs) + tuning.resolve_run_config(**inputs) assert not tuner.warmuped_configs_set @pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") @pytest.mark.parametrize( - "mode,width,projection_mode", + "mode,width,compute_mode", [ - ("gdn", 4, "inline"), + ("compact", 4, "inline"), ("replay", 4, "inline"), ("replay", 1, "inline"), ("replay", 4, "precompute"), @@ -149,17 +150,17 @@ def test_layout_is_fixed_before_graph_capture_and_cached_configs_do_not_execute( ) @pytest.mark.parametrize("dtype", [torch.float32, torch.bfloat16]) @pytest.mark.parametrize("kda", [False, True], ids=["gdn", "kda"]) -def test_joint_tuning_isolated_state_and_graph_replay(monkeypatch, mode, dtype, width, projection_mode, kda): - inputs = make_inputs(mode, dtype, "cuda", width=width, projection_mode=projection_mode, kda=kda) +def test_joint_tuning_isolated_state_and_graph_replay(monkeypatch, mode, dtype, width, compute_mode, kda): + inputs = make_inputs(mode, dtype, "cuda", width=width, compute_mode=compute_mode, kda=kda) if width == 1: for name in ("q", "k", "v"): inputs[name] = inputs[name].transpose(0, 1) inputs["cu_seqlens"] = None - cache = inputs["cache"] - for name, x in tensors(cache).items(): + ssm_backend = inputs["ssm_backend"] + for name, x in tensors(ssm_backend).items(): if name != "state": x.zero_() - before = {name: x.clone() for name, x in tensors(cache).items()} + before = {name: x.clone() for name, x in tensors(ssm_backend).items()} tuner = tuning.select_config configs = [{"BV": 16, "num_warps": 1, "num_stages": 1}, {"BV": 32, "num_warps": 2, "num_stages": 2}] monkeypatch.setattr(tuner, "configs_gen_func", lambda: configs) @@ -167,51 +168,56 @@ def test_joint_tuning_isolated_state_and_graph_replay(monkeypatch, mode, dtype, timings = [] def checked_bench(*args, **kwargs): - assert kwargs["cache"] is not cache + assert kwargs["ssm_backend"] is not ssm_backend + assert kwargs["ssm_backend"].state.shape[0] == 1 elapsed = bench(*args, **kwargs) assert math.isfinite(elapsed), f"failed {mode} {kwargs['run_config']}" + if mode == "replay": + # The measured cycle must reach a fold, not just empty-history updates. + scratch = kwargs["ssm_backend"] + assert scratch.cursors[0].item() >= 2 * scratch.capacity timings.append(elapsed) - for name, x in tensors(cache).items(): + for name, x in tensors(ssm_backend).items(): torch.testing.assert_close(x, before[name], atol=0, rtol=0) return elapsed monkeypatch.setattr(tuner, "_bench", checked_bench) args = [inputs[name] for name in ("q", "k", "v", "a", "b", "a_log", "bias")] - reqs = torch.full((3,), cache.hold, device="cuda", dtype=torch.int32) - cache.prepare_decode(reqs, inputs["cu_seqlens"]) + reqs = torch.full((3,), ssm_backend.hold_req_id, device="cuda", dtype=torch.int32) + ssm_backend.prepare_decode(reqs, inputs["cu_seqlens"]) with Autotuner.autotune_warmup(AutotuneKernelType.DECODE_ATTENTION): # Exercise the real forward hook with the HOLD-only graph warmup input. - cache.forward(0, *args, reqs, inputs["cu_seqlens"]) + ssm_backend.forward(0, *args, reqs, inputs["cu_seqlens"]) assert len(timings) == len(configs) - selected = cache.get_run_config(*args, inputs["cu_seqlens"]) - assert selected in configs + selected = ssm_backend.get_run_config(*args, inputs["cu_seqlens"]) + assert asdict(selected) in configs # An explicit identical layout supplies a reference for capture/replay and # alternating real/HOLD rows. It must never trigger its own tuning. - expected = make_inputs(mode, dtype, "cuda", width=width, projection_mode=projection_mode, kda=kda)["cache"] - expected.run_config = cache.run_config + expected = make_inputs(mode, dtype, "cuda", width=width, compute_mode=compute_mode, kda=kda)["ssm_backend"] + expected.run_config = ssm_backend.run_config expected._config_is_fixed = True expected.state.copy_(before["state"]) - accepted = torch.zeros(cache.state.shape[1], device="cuda", dtype=torch.int32) + accepted = torch.zeros(ssm_backend.state.shape[1], device="cuda", dtype=torch.int32) def step(target): target.prepare_decode(reqs, inputs["cu_seqlens"]) out = [ - target.forward(layer, *args, reqs, inputs["cu_seqlens"], run_config=selected) + target.forward(layer, *args, reqs, inputs["cu_seqlens"], run_config=selected)[0] for layer in range(target.state.shape[0]) ] - if mode == "gdn": + if mode == "compact": target.accept_updates(reqs, accepted, run_config=selected) else: target.accept_updates(reqs, accepted) return out - reqs.fill_(cache.hold) - step(cache) + reqs.fill_(ssm_backend.hold_req_id) + step(ssm_backend) torch.cuda.synchronize() graph = torch.cuda.CUDAGraph() with torch.cuda.graph(graph): - actual_out = step(cache) - reqs.copy_(torch.tensor([2, 0, cache.hold], device="cuda", dtype=torch.int32)) + actual_out = step(ssm_backend) + reqs.copy_(torch.tensor([2, 0, ssm_backend.hold_req_id], device="cuda", dtype=torch.int32)) for iteration in range(12): accepted[0] = iteration % width accepted[2] = (iteration + 1) % width @@ -219,29 +225,31 @@ def step(target): expected_out = step(expected) for actual, ref in zip(actual_out, expected_out): torch.testing.assert_close(actual, ref, atol=0, rtol=0) - torch.testing.assert_close(cache.state, expected.state, atol=0, rtol=0) + torch.testing.assert_close(ssm_backend.state, expected.state, atol=0, rtol=0) if mode == "replay": - cursor = cache.cursors.clone() - snapshot = cache.snapshot_accepted_state(0) - torch.testing.assert_close(cache.cursors, cursor, atol=0, rtol=0) - torch.testing.assert_close(snapshot, expected.snapshot_accepted_state(0), atol=0, rtol=0) - cache.merge_accepted_updates(reqs[:2]) + cursor = ssm_backend.cursors.clone() + snapshot = ssm_backend.get_accepted_state(0) + torch.testing.assert_close(ssm_backend.cursors, cursor, atol=0, rtol=0) + torch.testing.assert_close(snapshot, expected.get_accepted_state(0), atol=0, rtol=0) + ssm_backend.merge_accepted_updates(reqs[:2]) expected.merge_accepted_updates(reqs[:2]) - torch.testing.assert_close(cache.state, expected.state, atol=0, rtol=0) - torch.testing.assert_close(cache.state[:, cache.hold], before["state"][:, cache.hold], atol=0, rtol=0) + torch.testing.assert_close(ssm_backend.state, expected.state, atol=0, rtol=0) + torch.testing.assert_close( + ssm_backend.state[:, ssm_backend.hold_req_id], before["state"][:, ssm_backend.hold_req_id], atol=0, rtol=0 + ) assert len(timings) == len(configs) -@pytest.mark.parametrize("mode", ["gdn", "replay"]) +@pytest.mark.parametrize("mode", ["compact", "replay"]) def test_kda_rule_and_gate_bound_separate_tuning_cache(mode): inputs = make_inputs(mode, kda=True) key = tuning.select_config._static_key(**inputs) assert key != tuning.select_config._static_key(**make_inputs(mode)) _, rebuilt = tuning.rebuild_inputs(**inputs) - assert rebuilt["cache"].kda - assert rebuilt["cache"].lower_bound == inputs["cache"].lower_bound + assert rebuilt["ssm_backend"].kda + assert rebuilt["ssm_backend"].gate_lower_bound == inputs["ssm_backend"].gate_lower_bound assert key == tuning.select_config._static_key(**rebuilt) - inputs["cache"].lower_bound = -2.5 + inputs["ssm_backend"].gate_lower_bound = -2.5 assert key != tuning.select_config._static_key(**inputs) @@ -253,51 +261,51 @@ def test_production_tuning_cache_can_be_written(monkeypatch, tmp_path, tp, dtype monkeypatch.setattr(kernel_config_module, "get_current_device_name", lambda: "NVIDIA H100 80GB HBM3 MIG 3g.40gb") heads = 64 // tp - cache = SimpleNamespace( + ssm_backend = SimpleNamespace( state=SimpleNamespace(shape=(34, 1, heads, 128, 128), dtype=dtype), keys=SimpleNamespace(dtype=torch.bfloat16), verify_width=3, num_key_heads=heads, capacity=8, kda=kda, - lower_bound=-5.0, - projection_mode="precompute", + gate_lower_bound=-5.0, + compute_mode="precompute", ) packed = torch.empty(1, 3, heads * 128 * 3, dtype=torch.bfloat16) q, k, v = [x.view(1, 3, heads, 128) for x in packed.chunk(3, -1)] a = torch.empty(3, heads * 128 if kda else heads, dtype=q.dtype) b = torch.empty(3, heads, dtype=q.dtype) - key = tuning.static_key(cache, "replay", q, k, v, a, b, torch.tensor([0, 3])) + key = tuning.static_key(ssm_backend, "replay", q, k, v, a, b, torch.tensor([0, 3])) path = tmp_path / KernelConfigs.get_config_file_name(key) path.write_text("{}") assert path.is_file() def test_batch_configs_survive_other_warmups_and_preserve_padding(monkeypatch): - inputs = make_inputs("gdn") - cache = inputs["cache"] + inputs = make_inputs("compact") + ssm_backend = inputs["ssm_backend"] small = dict(inputs, **{name: inputs[name][:, :4] for name in ("q", "k", "v")}) small["a"], small["b"] = inputs["a"][:4], inputs["b"][:4] small["cu_seqlens"] = torch.tensor([0, 4, 4], dtype=torch.int32) selected = [] - def choose(cache, mode, q, *args): + def choose(ssm_backend, mode, q, *args): config = {"BV": 16 if q.shape[1] == 4 else 32, "num_warps": 2} selected.append(config) return config monkeypatch.setattr(tuning, "select_config", choose) with Autotuner.autotune_warmup(AutotuneKernelType.DECODE_ATTENTION): - large_config = tuning.configure_cache(**inputs) - small_config = tuning.configure_cache(**small) - assert tuning.configure_cache(**inputs) is large_config + large_config = tuning.resolve_run_config(**inputs) + small_config = tuning.resolve_run_config(**small) + assert tuning.resolve_run_config(**inputs) is large_config assert len(selected) == 2 and large_config != small_config - assert tuning.configure_cache(**small) is small_config - assert cache.run_config == {"BV": 8, "num_warps": 1} + assert tuning.resolve_run_config(**small) is small_config + assert ssm_backend.run_config == {"BV": 8, "num_warps": 1} _, rebuilt = tuning.rebuild_inputs(**small) - assert rebuilt["workload"][0].tolist() == [0, rebuilt["cache"].hold] + assert rebuilt["workload"][0].tolist() == [0, rebuilt["ssm_backend"].hold_req_id] assert rebuilt["cu_seqlens"].tolist() == [0, 4, 4] - assert tuning.run_key(**{k: rebuilt[k] for k in ("cache", "q", "cu_seqlens")}) == (4, 2) + assert tuning.run_key(**{k: rebuilt[k] for k in ("ssm_backend", "q", "cu_seqlens")}) == (4, 2) # A graph token bucket can also exceed the request-pool capacity. _, limited = tuning.rebuild_inputs(**dict(inputs, cu_seqlens=torch.tensor([0, 4, 8], dtype=torch.int32))) assert limited["workload"][0].tolist() == [0, 1] @@ -305,7 +313,7 @@ def choose(cache, mode, q, *args): @pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") -@pytest.mark.parametrize("mode,projection", [("gdn", "inline"), ("replay", "inline"), ("replay", "precompute")]) +@pytest.mark.parametrize("mode,projection", [("compact", "inline"), ("replay", "inline"), ("replay", "precompute")]) @pytest.mark.parametrize("dtype", [torch.float32, torch.bfloat16]) @pytest.mark.parametrize("kda", [False, True]) @pytest.mark.parametrize("dynamic", [False, True]) @@ -315,15 +323,15 @@ def test_two_verify_graphs_accept_with_their_own_configs(monkeypatch, mode, proj from lightllm.common.basemodel.batch_objs import PostLayerOutput, ModelMtpOutputCollector from lightllm.common.req_manager.linear_att import ReqManagerForMamba - inputs = make_inputs(mode, dtype, "cuda", projection_mode=projection, kda=kda) - cache = inputs["cache"] - initial = cache.state.clone() + inputs = make_inputs(mode, dtype, "cuda", compute_mode=projection, kda=kda) + ssm_backend = inputs["ssm_backend"] + initial = ssm_backend.state.clone() configs = [{"BV": 8, "num_warps": 1, "num_stages": 1}, {"BV": 32, "num_warps": 2, "num_stages": 1}] monkeypatch.setattr(tuning, "select_config", lambda *args: configs[args[2].shape[1] > 8]) manager = object.__new__(ReqManagerForMamba) - manager.recurrent_state = object.__new__(RecurrentState) - manager.recurrent_state.updates = cache - manager.recurrent_state.accepted_index = torch.zeros(9, dtype=torch.int32, device="cuda") + manager.linear_state_manager = object.__new__(LinearStateManager) + manager.linear_state_manager.ssm_backend = ssm_backend + manager.linear_state_manager.accepted_state_index = torch.zeros(9, dtype=torch.int32, device="cuda") buckets = [] for index, ids in enumerate(([0], [1, 3])): # The physical verify buckets include HOLD, but acceptance is unpadded. @@ -331,44 +339,52 @@ def test_two_verify_graphs_accept_with_their_own_configs(monkeypatch, mode, proj args = [inputs[name] for name in ("q", "k", "v", "a", "b", "a_log", "bias")] args[:3] = [x[:, : count * 4].contiguous() for x in args[:3]] args[3:5] = [x[: count * 4].contiguous() for x in args[3:5]] - capacity = min(count * 4, cache.hold) if dynamic else count - reqs = torch.tensor(ids + [cache.hold] * (capacity - len(ids)), dtype=torch.int32, device="cuda") + capacity = min(count * 4, ssm_backend.hold_req_id) if dynamic else count + reqs = torch.tensor(ids + [ssm_backend.hold_req_id] * (capacity - len(ids)), dtype=torch.int32, device="cuda") cu = (torch.arange(capacity + 1, dtype=torch.int32, device="cuda") * 4).clamp_max(count * 4) with Autotuner.autotune_warmup(AutotuneKernelType.DECODE_ATTENTION): - config = cache.get_run_config(*args, cu) - assert config == configs[index] - ref = make_inputs(mode, dtype, "cuda", projection_mode=projection, kda=kda)["cache"] + config = ssm_backend.get_run_config(*args, cu) + assert asdict(config) == configs[index] + ref = make_inputs(mode, dtype, "cuda", compute_mode=projection, kda=kda)["ssm_backend"] ref.state.copy_(initial) - ref.run_config = dict(config) + ref.run_config = asdict(config) ref._config_is_fixed = True def forward(target): target.prepare_decode(reqs, cu) - return [target.forward(layer, *args, reqs, cu, run_config=config) for layer in range(2)] - - forward(cache) + outputs = [] + selected = None + for layer in range(2): + result, actual_config = target.forward(layer, *args, reqs, cu, run_config=selected) + if selected is not None: + assert actual_config is selected + selected = actual_config + outputs.append(result) + return outputs, selected + + forward(ssm_backend) graph = torch.cuda.CUDAGraph() with torch.cuda.graph(graph): - outputs = forward(cache) + outputs, config = forward(ssm_backend) state = SimpleNamespace( hidden_collector=SimpleNamespace(finish_output=lambda **kw: ModelMtpOutputCollector()), prompt_logics=None, - ssm_run_config=SSMRunConfig.from_mapping(config), - logits_row_stride=1, + ssm_run_config=config, + input_rows_per_logit=1, ) output = TpPartBaseModel._create_model_output( None, PostLayerOutput(logits=outputs[0].view(count * 4, -1)), state ) output = TpPartBaseModel._create_unpad_decode_model_output(None, output, len(ids) * 4) - assert dict(output.ssm_run_config) == config + assert output.ssm_run_config is config buckets.append((graph, outputs, output, ref, args, reqs, cu, ids)) # Published contexts and captured launches must survive later changes to # the tuner's dictionaries, including when two graphs replay alternately. for config in configs: config.update(BV=128, num_warps=4, num_stages=3) - cache.state.copy_(initial) + ssm_backend.state.copy_(initial) if mode == "replay": - cache.cursors.zero_() + ssm_backend.cursors.zero_() req_order = [0, 1, 3] rows = torch.tensor([r for r in req_order for _ in range(4)], dtype=torch.int32, device="cuda") mtp = torch.arange(4, dtype=torch.int32, device="cuda").repeat(3) @@ -380,7 +396,7 @@ def forward(target): graph.replay() ref.prepare_decode(reqs, cu) for layer in range(2): - expected = ref.forward(layer, *args, reqs, cu) + expected = ref.forward(layer, *args, reqs, cu)[0] torch.testing.assert_close(outputs[layer], expected, atol=0, rtol=0) counts = [1 + (iteration + i) % 4 for i in range(3)] flags = torch.tensor([int(j < n) for n in counts for j in range(4)], dtype=torch.int32, device="cuda") @@ -390,11 +406,11 @@ def forward(target): mtp, flags, 4, - recurrent_state_batches=tuple((len(b[-1]), b[2].ssm_run_config) for b in buckets), + linear_state_batches=tuple((len(b[-1]), b[2].ssm_run_config) for b in buckets), ) for _, _, _, ref, _, reqs, _, ids in buckets: - ref.accept_updates(reqs[:-1], manager.recurrent_state.accepted_index) - torch.testing.assert_close(cache.state[:, ids], ref.state[:, ids], atol=0, rtol=0) + ref.accept_updates(reqs[:-1], manager.linear_state_manager.accepted_state_index) + torch.testing.assert_close(ssm_backend.state[:, ids], ref.state[:, ids], atol=0, rtol=0) if mode == "replay": - assert torch.equal(cache.cursors[ids], ref.cursors[ids]) - assert len(cache.batch_configs) == 2 + assert torch.equal(ssm_backend.cursors[ids], ref.cursors[ids]) + assert len(ssm_backend.batch_configs) == 2 diff --git a/unit_tests/common/basemodel/triton_kernel/linear_att/test_ssm_autotune_memory.py b/unit_tests/common/basemodel/triton_kernel/linear_att/test_ssm_autotune_memory.py index 022c71227b..398100ce7c 100644 --- a/unit_tests/common/basemodel/triton_kernel/linear_att/test_ssm_autotune_memory.py +++ b/unit_tests/common/basemodel/triton_kernel/linear_att/test_ssm_autotune_memory.py @@ -2,20 +2,22 @@ import torch from lightllm.common.basemodel.triton_kernel.linear_att import ssm_autotune as tuning -from lightllm.common.basemodel.triton_kernel.linear_att.replayssm import ReplaySSMCache +from lightllm.common.basemodel.triton_kernel.linear_att.replayssm import ReplaySSMBackend @pytest.mark.parametrize("width", [1, 5]) -def test_full_layer_scratch_exceeds_h200_headroom_without_allocating(width): +def test_scratch_budget_is_independent_of_model_layer_count(width): state = torch.empty((48, 257, 48, 128, 128), dtype=torch.bfloat16, device="meta") - cache = ReplaySSMCache(state, 16, width, num_key_heads=16) + ssm_backend = ReplaySSMBackend(state, 16, width, num_key_heads=16) tokens = 256 * width inputs = [ torch.empty(shape, dtype=torch.bfloat16, device="meta") for shape in [(tokens, 16, 128), (tokens, 16, 128), (tokens, 48, 128), (tokens, 48), (tokens, 48), (48,), (48,)] ] - required = tuning.scratch_size_bytes(cache, 256, 256, inputs) - assert 36 * 1024 ** 3 < required < 37 * 1024 ** 3 + required = tuning.scratch_size_bytes(ssm_backend, 256, 256, inputs) + assert required < 1024 ** 3 + single_layer = ReplaySSMBackend(state[:1], 16, width, num_key_heads=16) + assert required == tuning.scratch_size_bytes(single_layer, 256, 256, inputs) @pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") @@ -25,8 +27,8 @@ def test_tuning_budget_falls_back_without_touching_serving_state(monkeypatch, tm from lightllm.common.triton_utils.autotuner import Autotuner, AutotuneKernelType state = torch.zeros((2, 4, 2, 32, 32), dtype=torch.bfloat16, device="cuda") - cache = ReplaySSMCache(state, 8, width, num_key_heads=2) - for tensor in vars(cache).values(): + ssm_backend = ReplaySSMBackend(state, 8, width, num_key_heads=2) + for tensor in vars(ssm_backend).values(): if isinstance(tensor, torch.Tensor): tensor.zero_() tokens = 3 * width @@ -37,8 +39,10 @@ def test_tuning_budget_falls_back_without_touching_serving_state(monkeypatch, tm cu = torch.arange(4, device="cuda", dtype=torch.int32) * width if width > 1 else None if cu is not None: q = q.unsqueeze(0) - inputs = dict(cache=cache, mode="replay", q=q, k=q, v=q, a=gates, b=gates, a_log=a_log, bias=bias, cu_seqlens=cu) - before = {name: x.clone() for name, x in vars(cache).items() if isinstance(x, torch.Tensor)} + inputs = dict( + ssm_backend=ssm_backend, mode="replay", q=q, k=q, v=q, a=gates, b=gates, a_log=a_log, bias=bias, cu_seqlens=cu + ) + before = {name: x.clone() for name, x in vars(ssm_backend).items() if isinstance(x, torch.Tensor)} monkeypatch.setattr(torch.cuda, "mem_get_info", lambda device: (0, 1024 ** 3)) monkeypatch.setattr(torch.cuda, "memory_reserved", lambda device: 0) monkeypatch.setattr(torch.cuda, "memory_allocated", lambda device: 0) @@ -52,14 +56,14 @@ def test_tuning_budget_falls_back_without_touching_serving_state(monkeypatch, tm tuner, "_bench", lambda *args, **kwargs: pytest.fail("Insufficient memory must skip benchmarks") ) with Autotuner.autotune_warmup(AutotuneKernelType.DECODE_ATTENTION): - config = tuning.configure_cache(**inputs) - assert tuning.configure_cache(**inputs) is config + config = tuning.resolve_run_config(**inputs) + assert tuning.resolve_run_config(**inputs) is config assert config == {"BV": 32, "num_warps": 1} assert not list(tmp_path.iterdir()) for name, expected in before.items(): - torch.testing.assert_close(getattr(cache, name), expected, atol=0, rtol=0) + torch.testing.assert_close(getattr(ssm_backend, name), expected, atol=0, rtol=0) reqs = torch.arange(3, device="cuda", dtype=torch.int32) - cache.prepare_decode(reqs, cu) - cache.forward(0, q, q, q, gates, gates, a_log, bias, reqs, cu, run_config=config) + ssm_backend.prepare_decode(reqs, cu) + ssm_backend.forward(0, q, q, q, gates, gates, a_log, bias, reqs, cu, run_config=config) torch.cuda.synchronize() diff --git a/unit_tests/common/basemodel/triton_kernel/test_greedy_sampling.py b/unit_tests/common/basemodel/triton_kernel/test_greedy_sampling.py deleted file mode 100644 index 7ee34afe72..0000000000 --- a/unit_tests/common/basemodel/triton_kernel/test_greedy_sampling.py +++ /dev/null @@ -1,143 +0,0 @@ -import pytest -import torch - -from lightllm.common.basemodel.triton_kernel.post_process.argmax import argmax -from lightllm.common.basemodel.triton_kernel.post_process.greedy import greedy_sample - -pytestmark = pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") - - -@pytest.mark.parametrize("vocab", [13, 4101, 248320]) -def test_greedy_matches_torch_with_special_logits(vocab): - torch.manual_seed(417) - logits = torch.randn(7, vocab, device="cuda") - logits[0].fill_(-float("inf")) - logits[1, 2] = float("nan") - logits[1, 9] = float("nan") - logits[2, 4] = float("inf") - logits[3, 4] = logits[3, 9] = 100 - temperatures = torch.tensor([1, 1, 1, 0.5, 0.7, 1.3, 2.0], device="cuda") - expected = logits / temperatures[:, None] - ids = expected.argmax(-1) - probs = expected.softmax(-1).gather(1, ids[:, None]).log().flatten() - actual_ids, actual_probs = greedy_sample(logits, temperatures) - assert torch.equal(ids, actual_ids) - assert torch.equal(argmax(expected), ids) - torch.testing.assert_close(logits, expected, atol=0, rtol=0, equal_nan=True) - torch.testing.assert_close(actual_probs, probs, atol=3e-6, rtol=1e-6, equal_nan=True) - - -@pytest.mark.parametrize("vocab", [13, 4101]) -def test_identity_penalty_skip_preserves_eos_and_masks(vocab): - from types import SimpleNamespace - from lightllm.common.basemodel.triton_kernel.post_process.apply_penalty_gpu_cache import apply_penalty_gpu_cache - - torch.manual_seed(41) - manager = SimpleNamespace( - vocab_size=vocab, - req_to_presence_penalty=torch.zeros(7, device="cuda"), - req_to_frequency_penalty=torch.zeros(7, device="cuda"), - req_to_repetition_penalty=torch.ones(7, device="cuda"), - req_to_out_token_id_counter=torch.randint(0, 20, (7, vocab), device="cuda", dtype=torch.int32), - req_to_exponential_decay_length_penalty=torch.tensor([1.0, 1.01, 1.2, 1.0, 1.5, 0.9, 1.0], device="cuda"), - ) - logits = torch.randn(5, vocab, device="cuda") - original = logits.clone() - params = dict( - b_req_idx=torch.tensor([5, 3, 0, 4, 1], device="cuda", dtype=torch.int32), - b_length_penalty_param=torch.tensor([8, 0, 3, 4, 2], device="cuda", dtype=torch.int32), - b_mask_eos_reqs=torch.tensor([False, True, False, True, False], device="cuda"), - eos_ids=torch.tensor([2, 9], device="cuda", dtype=torch.int32), - sampling_params_manager=manager, - ) - apply_penalty_gpu_cache(original, **params) - apply_penalty_gpu_cache(logits, **params, skip_token_penalties=True) - torch.testing.assert_close(logits, original, rtol=0, atol=0) - assert (logits[[1, 3]][:, [2, 9]] == -10000000).all() - - -@pytest.mark.parametrize("greedy", [True, False]) -@pytest.mark.parametrize("identity", [True, False]) -def test_full_sampler_with_penalties_invalid_ids_and_random_sampling(monkeypatch, greedy, identity): - from types import SimpleNamespace - import lightllm.server.router.model_infer.mode_backend.generic_post_process as post - - vocab, batch = 4101, 4 - manager = SimpleNamespace( - penalty_counter_mode="gpu_counter", - req_to_bitmask_enabled=None, - vocab_size=vocab, - req_to_presence_penalty=torch.full((batch,), 0.0 if identity else 0.5, device="cuda"), - req_to_frequency_penalty=torch.full((batch,), 0.0 if identity else 0.1, device="cuda"), - req_to_repetition_penalty=torch.full((batch,), 1.0 if identity else 1.2, device="cuda"), - req_to_out_token_id_counter=torch.randint(0, 10, (batch, vocab), device="cuda", dtype=torch.int32), - req_to_exponential_decay_length_penalty=torch.full((batch,), 1.01, device="cuda"), - ) - reqs = [] - for index in range(batch): - param = SimpleNamespace( - presence_penalty=0.0 if identity else 0.5, - frequency_penalty=0.0 if identity else 0.1, - repetition_penalty=1.0 if identity else 1.2, - temperature=0.7 + 0.2 * index, - top_k=1 if greedy else vocab, - top_p=1.0, - min_new_tokens=12 if index % 2 else 0, - exponential_decay_length_penalty=SimpleNamespace(to_tuple=lambda: (0, 1.01)), - ) - reqs.append( - SimpleNamespace( - sampling_param=SimpleNamespace(shm_param=param, invalid_token_ids=[0, 5, 13]), - shm_req=SimpleNamespace(input_len=3), - get_cur_total_len=lambda: 10, - vocab_size=vocab, - generator=None, - req_idx=index, - ) - ) - monkeypatch.setattr( - post, "g_infer_context", SimpleNamespace(req_manager=SimpleNamespace(req_sampling_params_manager=manager)) - ) - monkeypatch.setattr( - post, - "g_pin_mem_manager", - SimpleNamespace(gen_from_list=lambda key, data, dtype: torch.tensor(data, dtype=dtype)), - ) - torch.manual_seed(62) - logits = torch.randn(batch, vocab, device="cuda") - torch.manual_seed(31) - actual = post.sample(logits.clone(), reqs, eos_id=[2, 9]) - - original_penalty = post.apply_penalty_gpu_cache - - def unskipped_penalty(**kwargs): - kwargs["skip_token_penalties"] = False - return original_penalty(**kwargs) - - def reference_greedy(logits, temperatures): - logits.div_(temperatures[:, None]) - probs = logits.softmax(-1) - ids = logits.argmax(-1) - return ids, probs.gather(1, ids[:, None]).log().flatten() - - monkeypatch.setattr(post, "apply_penalty_gpu_cache", unskipped_penalty) - monkeypatch.setattr(post, "greedy_sample", reference_greedy) - torch.manual_seed(31) - expected = post.sample(logits.clone(), reqs, eos_id=[2, 9]) - torch.testing.assert_close(actual[0], expected[0], rtol=0, atol=0) - torch.testing.assert_close(actual[1], expected[1], rtol=1e-6, atol=3e-6) - - -def test_greedy_fully_masked_partition_keeps_finite_logprob(): - logits = torch.randn(3, 8201, device="cuda") - logits[0, :4096] = -float("inf") - logits[1, 4096:8192] = -float("inf") - logits[2, 8192:] = -float("inf") - temperatures = torch.tensor([0.7, 1.0, 1.3], device="cuda") - scaled = logits / temperatures[:, None] - expected_ids = scaled.argmax(-1) - expected_logp = scaled.softmax(-1).gather(1, expected_ids[:, None]).log().flatten() - ids, logp = greedy_sample(logits, temperatures) - torch.testing.assert_close(ids, expected_ids, rtol=0, atol=0) - torch.testing.assert_close(logp, expected_logp, rtol=1e-6, atol=3e-6) - assert torch.isfinite(logp).all() diff --git a/unit_tests/common/test_recurrent_state.py b/unit_tests/common/test_linear_state.py similarity index 77% rename from unit_tests/common/test_recurrent_state.py rename to unit_tests/common/test_linear_state.py index 227548e9b6..3c0aa0c5df 100644 --- a/unit_tests/common/test_recurrent_state.py +++ b/unit_tests/common/test_linear_state.py @@ -4,7 +4,7 @@ import torch from lightllm.common.kv_cache_mem_manager.qwen3next_mem_manager import Qwen3NextLinearAttPageHelper -from lightllm.common.recurrent_state import RecurrentState +from lightllm.common.linear_state import LinearStateManager from lightllm.common.req_manager.linear_att import ReqManagerForMamba from lightllm.common.state_cache_manager import LinearAttCacheConfig, LinearAttCacheManager @@ -35,32 +35,34 @@ def test_checkpoint_and_pd_export_same_accepted_prefix_and_restore_reused_slot(m all_layer_num=3, ) args = SimpleNamespace(mtp_step=step, ssm_state_mode=mode, ssm_replay_config={"cache_len": 4}) - state = RecurrentState(3, config, args) + state = LinearStateManager(3, config, args) manager = object.__new__(ReqManagerForMamba) - manager.recurrent_state = state + manager.linear_state_manager = state cache = LinearAttCacheManager(3, config) manager.mem_manager = SimpleNamespace(big_page_buffers=cache) - mem = SimpleNamespace(linear_config=config, recurrent_state=state) + mem = SimpleNamespace(linear_config=config, linear_state_manager=state) helper = Qwen3NextLinearAttPageHelper(mem) src, dst, offset = 1, 2, min(1, step) width = config.get_mtp_conv_state_shape(step)[-1] state.conv_state.buffer.copy_(torch.arange(width, device="cuda").to(torch.bfloat16)) for slot in range(state.ssm_state.buffer.shape[1]): state.ssm_state.buffer[:, slot].fill_(slot + 1) - if state.accepted_index is not None: - state.accepted_index[src] = offset + if state.accepted_state_index is not None: + state.accepted_state_index[src] = offset expected_conv = state.conv_state.buffer[:, src, ..., offset : offset + 3].clone() - expected_ssm = state.ssm_state.buffer[:, src * state.slots_per_req + (offset if mode == "native" else 0)].clone() + expected_ssm = state.ssm_state.buffer[ + :, src * state.ssm_slots_per_req + (offset if mode == "native" else 0) + ].clone() if mode == "replay": # One pending accepted update: beta=0 leaves only decay of the base. - state.updates.raw_keys.zero_() - state.updates.raw_values.zero_() - state.updates.gates.fill_(-0.5) - state.updates.betas.zero_() - state.updates.cursors[src] = 1 + state.ssm_backend.raw_keys.zero_() + state.ssm_backend.raw_values.zero_() + state.ssm_backend.gates.fill_(-0.5) + state.ssm_backend.betas.zero_() + state.ssm_backend.cursors[src] = 1 expected_ssm = (expected_ssm.float() * torch.exp(torch.tensor(-0.5))).to(dtype) before = state.ssm_state.buffer.clone() - snapshot = state.snapshot(src) + snapshot = state.get_accepted_state(src) torch.testing.assert_close(snapshot.ssm, expected_ssm, rtol=2e-5, atol=2e-5) expected_ssm = snapshot.ssm.clone() @@ -78,13 +80,13 @@ def test_checkpoint_and_pd_export_same_accepted_prefix_and_restore_reused_slot(m torch.testing.assert_close(ssm_page, expected_ssm, rtol=0, atol=0) assert torch.equal(state.ssm_state.buffer, before) if mode == "replay": - assert state.updates.cursors[src].item() == 1 + assert state.ssm_backend.cursors[src].item() == 1 def dirty_destination(): - if state.accepted_index is not None: - state.accepted_index[dst] = step + if state.accepted_state_index is not None: + state.accepted_state_index[dst] = step if mode == "replay": - state.updates.cursors[dst] = 3 + state.ssm_backend.cursors[dst] = 3 for restore in ( lambda: manager.restore_state(SimpleNamespace(req_idx=dst), cache, 0), @@ -92,13 +94,13 @@ def dirty_destination(): ): dirty_destination() restore() - actual = state.snapshot(dst) + actual = state.get_accepted_state(dst) torch.testing.assert_close(actual.conv, expected_conv, rtol=0, atol=0) torch.testing.assert_close(actual.ssm, expected_ssm, rtol=0, atol=0) - if state.accepted_index is not None: - assert state.accepted_index[dst].item() == 0 + if state.accepted_state_index is not None: + assert state.accepted_state_index[dst].item() == 0 if mode == "replay": - assert state.updates.cursors[dst].item() == 0 + assert state.ssm_backend.cursors[dst].item() == 0 # A skipped checkpoint must not overwrite another cache owner. untouched = tuple(x.clone() for x in cache.get_state_cache(2)) @@ -108,5 +110,5 @@ def dirty_destination(): assert torch.equal(actual, expected) state.reset(dst) assert not state.conv_state.buffer[:, dst].any() - start = dst * state.slots_per_req - assert not state.ssm_state.buffer[:, start : start + state.slots_per_req].any() + start = dst * state.ssm_slots_per_req + assert not state.ssm_state.buffer[:, start : start + state.ssm_slots_per_req].any() diff --git a/unit_tests/common/test_recurrent_state_contract.py b/unit_tests/common/test_linear_state_contract.py similarity index 66% rename from unit_tests/common/test_recurrent_state_contract.py rename to unit_tests/common/test_linear_state_contract.py index b3d36a99ab..22ac52b792 100644 --- a/unit_tests/common/test_recurrent_state_contract.py +++ b/unit_tests/common/test_linear_state_contract.py @@ -5,31 +5,27 @@ import pytest import torch -from lightllm.common.recurrent_state import RecurrentState -from lightllm.common.recurrent_state_types import SSMRunConfig +from lightllm.common.linear_state import LinearStateManager, SSMRunConfig from lightllm.common.state_cache_manager import Glm5NextCacheConfig, LinearAttCacheConfig from lightllm.server.core.objs.start_args_type import StartArgs def test_published_config_is_detached_immutable_and_pickleable(): tuning = {"BV": 8, "num_warps": 1} - config = SSMRunConfig.from_mapping(tuning) + config = SSMRunConfig(**tuning) tuning.update(BV=128, num_warps=4, num_stages=1) - assert dict(config) == {"BV": 8, "num_warps": 1, "num_stages": 3} + assert config == SSMRunConfig(BV=8, num_warps=1, num_stages=3) with pytest.raises(FrozenInstanceError): config.BV = 128 - with pytest.raises(TypeError): - config["BV"] = 128 restored = pickle.loads(pickle.dumps(config)) assert restored == config - assert dict(restored) == dict(config) def test_config_rejects_unrepresented_kernel_options(): - with pytest.raises(ValueError, match="Unknown SSM kernel configuration fields"): - SSMRunConfig.from_mapping({"BV": 32, "num_warps": 1, "future_layout_option": 2}) + with pytest.raises(TypeError, match="unexpected keyword argument"): + SSMRunConfig(**{"BV": 32, "num_warps": 1, "future_layout_option": 2}) with pytest.raises(TypeError, match="must be an integer"): - SSMRunConfig.from_mapping({"BV": [32], "num_warps": 1}) + SSMRunConfig(BV=[32], num_warps=1) @pytest.fixture @@ -57,14 +53,14 @@ def linear_config(): @pytest.mark.parametrize("layer,row", [(0, 0), (1, 1), (3, 2), (4, 3), (6, 4), (7, 5)]) def test_layer_buffers_follow_packed_linear_layers(linear_config, layer, row): - owner = object.__new__(RecurrentState) + owner = object.__new__(LinearStateManager) owner.linear_config = linear_config conv = torch.arange(6).view(6, 1) ssm = conv + 10 owner.conv_state = SimpleNamespace(buffer=conv) owner.ssm_state = SimpleNamespace(buffer=ssm) assert linear_config.get_linear_layer_index(layer) == row - actual_conv, actual_ssm = owner.layer_buffers(layer) + actual_conv, actual_ssm = owner.get_layer_buffers(layer) assert actual_conv.item() == row assert actual_ssm.item() == row + 10 @@ -98,8 +94,27 @@ def test_model_configuration_retains_nondefault_kda_gate_bound(): @pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") @pytest.mark.parametrize("mode", ["compact", "replay"]) -def test_kda_updater_receives_gate_bound_at_construction(linear_config, mode): +def test_kda_backend_receives_gate_bound_at_construction(linear_config, mode): linear_config.gate_lower_bound = -2.5 args = SimpleNamespace(mtp_step=2, ssm_state_mode=mode, ssm_replay_config={}) - owner = RecurrentState(2, linear_config, args, kda=True) - assert owner.updates.lower_bound == -2.5 + owner = LinearStateManager(2, linear_config, args, kda=True) + assert owner.ssm_backend.gate_lower_bound == -2.5 + + +@pytest.mark.parametrize("mode", ["compact", "replay", "parallel"]) +def test_ssm_backend_publishes_detached_run_config(mode): + from lightllm.common.basemodel.triton_kernel.linear_att.replayssm import ReplaySSMBackend + from lightllm.common.basemodel.triton_kernel.linear_att.compact_ssm import CompactSSMBackend + + ssm_backend = object.__new__(CompactSSMBackend if mode == "compact" else ReplaySSMBackend) + tuning = {"BV": 32, "num_warps": 2} + ssm_backend.run_config = tuning + ssm_backend._config_is_fixed = True + ssm_backend.compute_mode = "parallel" if mode == "parallel" else "inline" + config = ssm_backend.get_run_config(*([None] * 7)) + assert isinstance(config, SSMRunConfig) + tuning["BV"] = 128 + assert config.BV == 32 + assert config.num_stages == 3 + with pytest.raises(FrozenInstanceError): + config.BV = 128 diff --git a/unit_tests/models/test_post_layer_output.py b/unit_tests/models/test_post_layer_output.py index c8e081862d..9e56e47055 100644 --- a/unit_tests/models/test_post_layer_output.py +++ b/unit_tests/models/test_post_layer_output.py @@ -60,7 +60,7 @@ def candidates(*args): args = SimpleNamespace(target_vocab_topk_sampling=None, draft_vocab_topk_sampling=None) setattr(args, "draft_vocab_topk_sampling" if draft else "target_vocab_topk_sampling", 16 if enabled else None) monkeypatch.setattr(llama_post, "get_env_start_args", lambda: args) - state = SimpleNamespace(is_draft_model=draft, dist_group=None, b_logits_rows=None) + state = SimpleNamespace(is_draft_model=draft, dist_group=None, b_logits_row_idx=None) weight = SimpleNamespace(lm_head_weight_=Weight()) method = head._draft_lm_head_and_gather if draft else head._target_lm_head_and_gather output = method(torch.ones(3, 2), 3, weight, state) @@ -291,7 +291,7 @@ def test_model_output_and_unpadding_preserve_candidate_mapping() -> None: logits = torch.randn(3, 2) ids = torch.tensor([[10, 20], [30, 40], [50, 60]]) state = SimpleNamespace( - hidden_collector=NoopHiddenCollector(), prompt_logics=None, ssm_run_config=None, logits_row_stride=1 + hidden_collector=NoopHiddenCollector(), prompt_logics=None, ssm_run_config=None, input_rows_per_logit=1 ) output = model._create_model_output(PostLayerOutput(logits, ids), state) unpadded = model._create_unpad_decode_model_output(output, 2) diff --git a/unit_tests/server/router/model_infer/mtp_speculative/test_eagle_with_att.py b/unit_tests/server/router/model_infer/mtp_speculative/test_eagle_with_att.py index 94be2ef52a..06da6896b4 100644 --- a/unit_tests/server/router/model_infer/mtp_speculative/test_eagle_with_att.py +++ b/unit_tests/server/router/model_infer/mtp_speculative/test_eagle_with_att.py @@ -23,7 +23,7 @@ def forward(model_input): draft_model = SimpleNamespace( forward=forward, - use_replayssm_draft_tail=lambda: False, + should_project_draft_tail_logits=lambda: False, map_draft_vocab_to_main_vocab=lambda token_ids: token_ids + 100, ) proposer = Eagle3Proposer( @@ -104,7 +104,7 @@ def forward(model_input): proposer = EagleWithAttProposer( backend=SimpleNamespace( - draft_models=[SimpleNamespace(forward=forward, use_replayssm_draft_tail=lambda: False)] + draft_models=[SimpleNamespace(forward=forward, should_project_draft_tail_logits=lambda: False)] ), enable_dynmaic_mtp=False, ) @@ -161,7 +161,7 @@ def forward(model_input): ) backend = SimpleNamespace( - draft_models=[SimpleNamespace(forward=forward, use_replayssm_draft_tail=lambda: False)], + draft_models=[SimpleNamespace(forward=forward, should_project_draft_tail_logits=lambda: False)], _gen_argmax_token_ids=lambda output: output.logits[:, 0].long(), _gen_argmax_token_ids_and_prob=lambda output: ( output.logits[:, 0].long(), @@ -225,19 +225,19 @@ def test_selected_tail_logits_keep_complete_draft_kv_input(capacity): input_ids=None, b_position_delta=torch.zeros(10, dtype=torch.int32), mtp_draft_input_hiddens=None, - b_logits_rows=None, + b_logits_row_idx=None, ) def forward(model_input): assert model_input.batch_size == 10 assert model_input.mtp_draft_input_hiddens is hidden assert model_input.input_ids.numel() == 10 - torch.testing.assert_close(model_input.b_logits_rows, torch.tensor([1, 8])) - return ModelOutput(logits=torch.tensor([[7.0], [9.0]]), logits_row_stride=5) + torch.testing.assert_close(model_input.b_logits_row_idx, torch.tensor([1, 8])) + return ModelOutput(logits=torch.tensor([[7.0], [9.0]]), input_rows_per_logit=5) draft_model = SimpleNamespace( forward=forward, - use_replayssm_draft_tail=lambda: True, + should_project_draft_tail_logits=lambda: True, args=SimpleNamespace(running_max_req_size=capacity), ) proposer = EagleWithAttProposer( @@ -253,5 +253,5 @@ def forward(model_input): torch.tensor([2, 4]), ) torch.testing.assert_close(output.token_ids, torch.tensor([[7], [9]])) - assert target_input.b_logits_rows is None + assert target_input.b_logits_row_idx is None assert target_input.input_ids is None diff --git a/unit_tests/server/test_ssm_state_mode.py b/unit_tests/server/test_ssm_state_mode.py index 58dee5938c..57151db0ae 100644 --- a/unit_tests/server/test_ssm_state_mode.py +++ b/unit_tests/server/test_ssm_state_mode.py @@ -3,7 +3,7 @@ import pytest from lightllm.server.api_cli import make_argument_parser -from lightllm.server.api_start import _validate_ssm_state_mode +from lightllm.server.api_start import _validate_and_normalize_ssm_config from lightllm.server.core.objs.start_args_type import StartArgs @@ -49,9 +49,9 @@ def test_ssm_mode_validation(monkeypatch, mode, dtype, step, model, capacity, er ) if error: with pytest.raises(AssertionError, match=error): - _validate_ssm_state_mode(args) + _validate_and_normalize_ssm_config(args) else: - _validate_ssm_state_mode(args) + _validate_and_normalize_ssm_config(args) @pytest.mark.parametrize("model", ["qwen3_5", "glm5_next"]) @@ -61,10 +61,10 @@ def test_gdn_kda_mtp_width_limit(monkeypatch, model, mode, step): monkeypatch.setattr("lightllm.utils.config_utils.get_model_type", lambda _: model) args = StartArgs(ssm_state_mode=mode, mtp_step=step) if step <= 15: - _validate_ssm_state_mode(args) + _validate_and_normalize_ssm_config(args) else: with pytest.raises(AssertionError, match="GDN/KDA requires mtp_step <= 15"): - _validate_ssm_state_mode(args) + _validate_and_normalize_ssm_config(args) @pytest.mark.parametrize("model", ["qwen3_5", "glm5_next"]) @@ -79,7 +79,7 @@ def test_replay_large_cache_does_not_bypass_mtp_width_limit(monkeypatch, model, mtp_step=16, ) with pytest.raises(AssertionError, match="GDN/KDA requires mtp_step <= 15"): - _validate_ssm_state_mode(args) + _validate_and_normalize_ssm_config(args) def test_ssm_mode_defaults_and_removed_flag(): @@ -107,17 +107,17 @@ def test_programmatic_replay_config_rejected_before_model_loading(monkeypatch, o for key, value in overrides.items(): setattr(args, key, value) with pytest.raises(AssertionError, match=error): - _validate_ssm_state_mode(args) + _validate_and_normalize_ssm_config(args) -def test_replay_projection_mode_requires_replay(monkeypatch): +def test_replay_compute_mode_requires_replay(monkeypatch): monkeypatch.setattr("lightllm.utils.config_utils.get_model_type", lambda _: "qwen3_5") parser = make_argument_parser() args = parser.parse_args(["--ssm_state_mode", "replay", "--ssm_replay_config", '{"compute_mode":"precompute"}']) - _validate_ssm_state_mode(args) + _validate_and_normalize_ssm_config(args) args.ssm_state_mode = "compact" with pytest.raises(AssertionError, match="requires ssm_state_mode=replay"): - _validate_ssm_state_mode(args) + _validate_and_normalize_ssm_config(args) @pytest.mark.parametrize("fold", ["raw", "derived"]) @@ -132,10 +132,10 @@ def test_parallel_replay_configuration(monkeypatch, fold): mtp_mode="eagle_with_att", linear_att_ssm_data_type="bfloat16", ) - _validate_ssm_state_mode(args) + _validate_and_normalize_ssm_config(args) args.mtp_step = 0 with pytest.raises(AssertionError, match="requires MTP"): - _validate_ssm_state_mode(args) + _validate_and_normalize_ssm_config(args) def test_derived_replay_requires_parallel_configuration(): @@ -143,7 +143,7 @@ def test_derived_replay_requires_parallel_configuration(): args = StartArgs(ssm_state_mode="replay", ssm_replay_config={"fold_method": "derived"}) with pytest.raises(AssertionError, match="requires parallel"): - _validate_ssm_state_mode(args) + _validate_and_normalize_ssm_config(args) @pytest.mark.parametrize( @@ -163,14 +163,14 @@ def test_derived_replay_requires_parallel_configuration(): def test_invalid_replay_config(config, error): args = StartArgs(ssm_state_mode="replay", ssm_replay_config=config) with pytest.raises(AssertionError, match=error): - _validate_ssm_state_mode(args) + _validate_and_normalize_ssm_config(args) @pytest.mark.parametrize("mode", ["native", "compact"]) def test_replay_config_rejects_other_modes(mode): args = StartArgs(ssm_state_mode=mode, ssm_replay_config={"cache_len": 16}) with pytest.raises(AssertionError, match="requires ssm_state_mode=replay"): - _validate_ssm_state_mode(args) + _validate_and_normalize_ssm_config(args) def test_partial_replay_config_defaults_and_serialization(monkeypatch): @@ -180,11 +180,11 @@ def test_partial_replay_config_defaults_and_serialization(monkeypatch): args = make_argument_parser().parse_args( ["--ssm_state_mode", "replay", "--ssm_replay_config", '{"compute_mode":"precompute"}'] ) - _validate_ssm_state_mode(args) + _validate_and_normalize_ssm_config(args) assert args.ssm_replay_config == {"cache_len": 16, "compute_mode": "precompute", "fold_method": "raw"} original = StartArgs(ssm_state_mode="replay", ssm_replay_config=args.ssm_replay_config) restored = StartArgs(**json.loads(json.dumps(asdict(original)))) - _validate_ssm_state_mode(restored) + _validate_and_normalize_ssm_config(restored) assert restored.ssm_replay_config == args.ssm_replay_config assert StartArgs().ssm_replay_config == {} From c3818f8910600e0526b9da936640491be5cf3548 Mon Sep 17 00:00:00 2001 From: wangzaijun Date: Thu, 22 Oct 2026 01:59:06 +0000 Subject: [PATCH 32/40] style: revert unrelated blank-line changes in ReplaySSM PR --- lightllm/common/basemodel/attention/linear/gdn.py | 1 + lightllm/common/basemodel/basemodel.py | 1 + lightllm/common/basemodel/triton_kernel/mtp_utils.py | 1 + lightllm/common/kv_cache_mem_manager/mem_manager.py | 1 + lightllm/common/kv_cache_mem_manager/qwen3next_mem_manager.py | 1 - lightllm/common/state_cache_manager/linear_att.py | 1 + lightllm/models/llama/layer_infer/post_layer_infer.py | 3 ++- .../models/qwen3next/layer_infer/transformer_layer_infer.py | 3 ++- lightllm/server/api_start.py | 1 - lightllm/server/router/model_infer/infer_batch.py | 1 + 10 files changed, 10 insertions(+), 4 deletions(-) diff --git a/lightllm/common/basemodel/attention/linear/gdn.py b/lightllm/common/basemodel/attention/linear/gdn.py index cdd6016a31..483a603b7f 100644 --- a/lightllm/common/basemodel/attention/linear/gdn.py +++ b/lightllm/common/basemodel/attention/linear/gdn.py @@ -39,6 +39,7 @@ 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 diff --git a/lightllm/common/basemodel/basemodel.py b/lightllm/common/basemodel/basemodel.py index ddf7486c2e..653b12de2f 100755 --- a/lightllm/common/basemodel/basemodel.py +++ b/lightllm/common/basemodel/basemodel.py @@ -772,6 +772,7 @@ def _decode( @final def _context_forward(self, infer_state: InferStateInfo): + input_embs = self.pre_infer.context_forward(infer_state.input_ids, infer_state, self.pre_post_weight) if self.args.enable_dp_prefill_balance: assert not self.args.enable_prefill_cudagraph, "not support now" diff --git a/lightllm/common/basemodel/triton_kernel/mtp_utils.py b/lightllm/common/basemodel/triton_kernel/mtp_utils.py index 7b7a16bfe2..c72f8191df 100644 --- a/lightllm/common/basemodel/triton_kernel/mtp_utils.py +++ b/lightllm/common/basemodel/triton_kernel/mtp_utils.py @@ -203,6 +203,7 @@ def _fwd_kernel_mtp_scatter_next_token_ids( HAS_NEXT_TOKEN_SCORES: tl.constexpr, BLOCK_SIZE: tl.constexpr, ): + cur_index = tl.program_id(0) req_start_loc = tl.load(b_req_mtp_start_loc + cur_index) accept_len = tl.load(mtp_accept_len + cur_index) diff --git a/lightllm/common/kv_cache_mem_manager/mem_manager.py b/lightllm/common/kv_cache_mem_manager/mem_manager.py index fd5d814fe9..f606ceb30b 100755 --- a/lightllm/common/kv_cache_mem_manager/mem_manager.py +++ b/lightllm/common/kv_cache_mem_manager/mem_manager.py @@ -30,6 +30,7 @@ class MemoryManager: + operator_class = NormalMemOperator def __init__(self, size, dtype, head_num, head_dim, layer_num, always_copy=False, mem_fraction=0.9): 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 2989c9a5cd..0320632015 100644 --- a/lightllm/common/kv_cache_mem_manager/qwen3next_mem_manager.py +++ b/lightllm/common/kv_cache_mem_manager/qwen3next_mem_manager.py @@ -2,7 +2,6 @@ import torch import triton - from lightllm.utils.log_utils import init_logger from lightllm.common.kv_cache_mem_manager.mem_manager import MemoryManager from lightllm.utils.envs_utils import get_env_start_args diff --git a/lightllm/common/state_cache_manager/linear_att.py b/lightllm/common/state_cache_manager/linear_att.py index 12d262e608..ec6cae7666 100644 --- a/lightllm/common/state_cache_manager/linear_att.py +++ b/lightllm/common/state_cache_manager/linear_att.py @@ -6,6 +6,7 @@ from lightllm.utils.envs_utils import get_added_mtp_kv_layer_num, get_env_start_args from lightllm.utils.log_utils import init_logger from lightllm.utils.torch_dtype_utils import get_torch_dtype + from .base import StateCacheManager from .layer_cache import LayerCache diff --git a/lightllm/models/llama/layer_infer/post_layer_infer.py b/lightllm/models/llama/layer_infer/post_layer_infer.py index 02cb01ebc9..93fe21974a 100644 --- a/lightllm/models/llama/layer_infer/post_layer_infer.py +++ b/lightllm/models/llama/layer_infer/post_layer_infer.py @@ -1,5 +1,4 @@ import os - import torch import torch.functional as F import torch.distributed as dist @@ -256,6 +255,7 @@ def _draft_lm_head_and_gather( def token_forward( self, input_embdings: torch.Tensor, infer_state: LlamaInferStateInfo, layer_weight: LlamaPreAndPostLayerWeight ) -> PostLayerOutput: + return self._token_forward(input_embdings=input_embdings, infer_state=infer_state, layer_weight=layer_weight) def overlap_tpsp_token_forward( @@ -266,6 +266,7 @@ def overlap_tpsp_token_forward( infer_state1: LlamaInferStateInfo, layer_weight: BaseLayerWeight, ) -> tuple[PostLayerOutput, PostLayerOutput]: + output = self.token_forward(input_embdings, infer_state, layer_weight=layer_weight) output1 = self.token_forward(input_embdings1, infer_state1, layer_weight=layer_weight) diff --git a/lightllm/models/qwen3next/layer_infer/transformer_layer_infer.py b/lightllm/models/qwen3next/layer_infer/transformer_layer_infer.py index bb0f177f2e..4fe3e4e93c 100644 --- a/lightllm/models/qwen3next/layer_infer/transformer_layer_infer.py +++ b/lightllm/models/qwen3next/layer_infer/transformer_layer_infer.py @@ -2,8 +2,8 @@ from functools import partial import torch -import torch.distributed as dist +import torch.distributed as dist from lightllm.models.qwen3next.layer_weights.transformer_layer_weight import ( Qwen3NextTransformerLayerWeight, ) @@ -184,6 +184,7 @@ def _linear_prefill_cuda_graph_wrapper( infer_state: Qwen3NextInferStateInfo, layer_weight: Qwen3NextTransformerLayerWeight, ) -> Tuple[torch.Tensor, torch.Tensor]: + backend = infer_state.prefill_att_state1.backend mixed_qkvzba = mixed_qkvzba.contiguous() diff --git a/lightllm/server/api_start.py b/lightllm/server/api_start.py index 840537d794..52da4fb5e1 100644 --- a/lightllm/server/api_start.py +++ b/lightllm/server/api_start.py @@ -3,7 +3,6 @@ import uuid import subprocess import math - from lightllm.utils.start_utils import process_manager from .metrics.manager import start_metric_manager from .httpserver.req_id_generator_server import start_req_id_generator_server diff --git a/lightllm/server/router/model_infer/infer_batch.py b/lightllm/server/router/model_infer/infer_batch.py index 752492b646..f5d712a3fb 100644 --- a/lightllm/server/router/model_infer/infer_batch.py +++ b/lightllm/server/router/model_infer/infer_batch.py @@ -363,6 +363,7 @@ def filter_reqs(self, finished_reqs: List["InferReq"]): @torch.no_grad() def pause_reqs(self, pause_reqs: List["InferReq"], is_master_in_dp: bool): if pause_reqs: + free_token_index = [] for req in pause_reqs: if self.args.diverse_mode: From 15270021651c7a4799d42ca30a7d60bb99ea43c5 Mon Sep 17 00:00:00 2001 From: sufubao Date: Thu, 22 Oct 2026 14:06:14 +0800 Subject: [PATCH 33/40] refactor: keep SSM run configs in attention state --- .../common/basemodel/attention/base_att.py | 4 ++ .../common/basemodel/attention/linear/gdn.py | 15 +++-- .../common/basemodel/attention/linear/kda.py | 11 +++- lightllm/common/basemodel/basemodel.py | 9 ++- lightllm/common/basemodel/batch_objs.py | 6 +- lightllm/common/basemodel/cuda_graph.py | 16 +++++- lightllm/common/basemodel/infer_struct.py | 6 +- lightllm/common/linear_state.py | 20 +++++-- .../mode_backend/chunked_prefill/impl.py | 2 +- .../mode_backend/dp_backend/impl.py | 6 +- .../basemodel/test_cuda_graph_autotune.py | 57 +++++++++++++++++++ .../linear_att/test_ssm_autotune.py | 21 ++++--- .../common/test_linear_state_contract.py | 30 ++++++++++ unit_tests/models/test_post_layer_output.py | 6 +- 14 files changed, 173 insertions(+), 36 deletions(-) diff --git a/lightllm/common/basemodel/attention/base_att.py b/lightllm/common/basemodel/attention/base_att.py index 1122df3551..28db5bf474 100644 --- a/lightllm/common/basemodel/attention/base_att.py +++ b/lightllm/common/basemodel/attention/base_att.py @@ -10,6 +10,7 @@ from lightllm.utils.envs_utils import get_env_start_args if TYPE_CHECKING: + from lightllm.common.linear_state import LinearStateCommit from lightllm.common.basemodel.basemodel import TpPartBaseModel from lightllm.common.basemodel.infer_struct import InferStateInfo @@ -161,6 +162,9 @@ def prefill_att( @dataclass class BaseDecodeAttState(BaseAttState): + def get_linear_state_commit(self) -> Optional["LinearStateCommit"]: + return None + def copy_for_decode_cuda_graph(self, new_state: "BaseDecodeAttState"): """Copy tensor metadata into the buffers owned by the captured graph.""" for attr_name, attr_value in vars(new_state).items(): diff --git a/lightllm/common/basemodel/attention/linear/gdn.py b/lightllm/common/basemodel/attention/linear/gdn.py index 483a603b7f..7c6d2939a2 100644 --- a/lightllm/common/basemodel/attention/linear/gdn.py +++ b/lightllm/common/basemodel/attention/linear/gdn.py @@ -2,7 +2,7 @@ import dataclasses import weakref from abc import ABC, abstractmethod -from typing import TYPE_CHECKING +from typing import TYPE_CHECKING, Optional import torch @@ -22,6 +22,7 @@ ) if TYPE_CHECKING: + from lightllm.common.linear_state import LinearStateCommit, SSMRunConfig from lightllm.common.basemodel.basemodel import TpPartBaseModel from lightllm.common.basemodel.infer_struct import InferStateInfo from lightllm.models.qwen3next.infer_struct import Qwen3NextInferStateInfo @@ -201,12 +202,16 @@ def _gdn_prefill_kernel( @dataclasses.dataclass class GDNLinearAttDecodeAttState(BaseDecodeAttState): + run_config: Optional["SSMRunConfig"] = None backend: GDNLinearAttBackend = None 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 + def get_linear_state_commit(self) -> Optional["LinearStateCommit"]: + return self.infer_state.req_manager.linear_state_manager.create_commit(self.run_config) + 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: @@ -336,7 +341,7 @@ def _gdn_decode_kernel( linear_layer_idx = self.infer_state.req_manager.linear_config.get_linear_layer_index( layer_weight.layer_num_ ) - core_attn_out, self.infer_state.ssm_run_config = ssm_backend.forward( + core_attn_out, self.run_config = ssm_backend.forward( linear_layer_idx, query, key, @@ -346,7 +351,7 @@ def _gdn_decode_kernel( layer_weight.linear_A_log.weight, layer_weight.linear_dt_bias.weight, self.b_conv_buffer_idx, - run_config=self.infer_state.ssm_run_config, + run_config=self.run_config, ) return core_attn_out, z core_attn_out, _ = fused_recurrent_gated_delta_rule( @@ -399,7 +404,7 @@ def _gdn_mtp_kernel( linear_layer_idx = self.infer_state.req_manager.linear_config.get_linear_layer_index( layer_weight.layer_num_ ) - core_attn_out, self.infer_state.ssm_run_config = ssm_backend.forward( + core_attn_out, self.run_config = ssm_backend.forward( linear_layer_idx, query, key, @@ -410,7 +415,7 @@ def _gdn_mtp_kernel( layer_weight.linear_dt_bias.weight, self.b_conv_buffer_idx, cu_seqlens_q, - run_config=self.infer_state.ssm_run_config, + run_config=self.run_config, ) return core_attn_out assert self.b_ssm_buffer_idx.dim() == 2, "SSM buffer idx must be 2D [N, S+1]" diff --git a/lightllm/common/basemodel/attention/linear/kda.py b/lightllm/common/basemodel/attention/linear/kda.py index e28058826b..3e7a7634b6 100644 --- a/lightllm/common/basemodel/attention/linear/kda.py +++ b/lightllm/common/basemodel/attention/linear/kda.py @@ -6,7 +6,7 @@ import dataclasses import weakref -from typing import TYPE_CHECKING +from typing import TYPE_CHECKING, Optional import torch @@ -31,6 +31,7 @@ ) if TYPE_CHECKING: + from lightllm.common.linear_state import LinearStateCommit, SSMRunConfig from lightllm.common.basemodel.basemodel import TpPartBaseModel from lightllm.common.basemodel.infer_struct import InferStateInfo @@ -132,11 +133,15 @@ def prefill_att( @dataclasses.dataclass class KDALinearAttDecodeAttState(BaseDecodeAttState): + run_config: Optional["SSMRunConfig"] = None 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 + def get_linear_state_commit(self) -> Optional["LinearStateCommit"]: + return self.infer_state.req_manager.linear_state_manager.create_commit(self.run_config) + def init_state(self): mtp_step = self.backend.mtp_step if mtp_step == 0: @@ -228,7 +233,7 @@ def decode_att( ssm_backend = self.infer_state.req_manager.linear_state_manager.ssm_backend if ssm_backend is not None: linear_layer_idx = self.infer_state.req_manager.linear_config.get_linear_layer_index(layer_num) - output, self.infer_state.ssm_run_config = ssm_backend.forward( + output, self.run_config = ssm_backend.forward( linear_layer_idx, q, k, @@ -239,7 +244,7 @@ def decode_att( layer_weight.linear_dt_bias.weight, self.b_conv_buffer_idx, self.b1_mtp_cu_q_seq_len, - run_config=self.infer_state.ssm_run_config, + run_config=self.run_config, ) return output output, _ = fused_recurrent_kda( diff --git a/lightllm/common/basemodel/basemodel.py b/lightllm/common/basemodel/basemodel.py index 653b12de2f..1488937daf 100755 --- a/lightllm/common/basemodel/basemodel.py +++ b/lightllm/common/basemodel/basemodel.py @@ -618,12 +618,19 @@ def _create_padded_prefill_model_input(self, model_input: ModelInput, new_handle return new_model_input def _create_model_output(self, post_output: PostLayerOutput, infer_state: InferStateInfo) -> ModelOutput: + linear_state_commit = None + for state in (infer_state.decode_att_state, infer_state.decode_att_state1): + if state is not None: + commit = state.get_linear_state_commit() + if commit is not None: + assert linear_state_commit is None, "multiple linear state commits in one forward" + linear_state_commit = commit output = ModelOutput( logits=post_output.logits.contiguous(), mtp_collector=infer_state.hidden_collector.finish_output(infer_state=infer_state), prompt_logics=infer_state.prompt_logics, logits_token_ids=post_output.logits_token_ids, - ssm_run_config=infer_state.ssm_run_config, + linear_state_commit=linear_state_commit, input_rows_per_logit=infer_state.input_rows_per_logit, ) return output diff --git a/lightllm/common/basemodel/batch_objs.py b/lightllm/common/basemodel/batch_objs.py index d2da5d377f..3b1a27ebc8 100644 --- a/lightllm/common/basemodel/batch_objs.py +++ b/lightllm/common/basemodel/batch_objs.py @@ -7,7 +7,7 @@ from lightllm.utils.tensor_utils import tensor_to_no_ref_tensor if TYPE_CHECKING: - from lightllm.common.linear_state import SSMRunConfig + from lightllm.common.linear_state import LinearStateCommit @dataclass @@ -210,8 +210,8 @@ class ModelOutput: # target 模型会将候选值回填到完整词表 logits，因此不需要该映射。 logits_token_ids: Optional[torch.Tensor] = None - # Immutable SSM configuration for this forward/graph; preserved when unpadding. - ssm_run_config: Optional["SSMRunConfig"] = None + # Detached state-commit metadata for this forward/graph; preserved when unpadding. + linear_state_commit: Optional["LinearStateCommit"] = None # Number of physical decode rows represented by each selected logits row. input_rows_per_logit: int = 1 diff --git a/lightllm/common/basemodel/cuda_graph.py b/lightllm/common/basemodel/cuda_graph.py index 5a38023b89..21cf83e5f4 100644 --- a/lightllm/common/basemodel/cuda_graph.py +++ b/lightllm/common/basemodel/cuda_graph.py @@ -2,6 +2,7 @@ import copy import bisect import math +import weakref from typing import Optional import torch @@ -23,6 +24,17 @@ logger = init_logger(__name__) +def _copy_warmup_state(infer_state: InferStateInfo) -> InferStateInfo: + warmup_state = copy.copy(infer_state) + for name in ("decode_att_state", "decode_att_state1"): + state = getattr(infer_state, name, None) + if state is not None: + state = copy.copy(state) + state._infer_state_weak_ref = weakref.ref(warmup_state) + setattr(warmup_state, name, state) + return warmup_state + + class CudaGraph: # CudaGraph forward pass for the decoding stage. @@ -125,7 +137,7 @@ def _capture_decode(self, decode_func, infer_state: InferStateInfo): torch.cuda.synchronize() # 在正式捕获前调优 decode attention，退出作用域后再捕获选定的配置。 with Autotuner.autotune_warmup(AutotuneKernelType.DECODE_ATTENTION): - decode_func(copy.copy(infer_state)) + decode_func(_copy_warmup_state(infer_state)) torch.cuda.synchronize() for param_name in set(vars(infer_state).keys()): if param_name not in pure_para_set: @@ -158,7 +170,7 @@ def _capture_decode_overlap( pure_para_set1 = set(vars(infer_state1).keys()) torch.cuda.synchronize() with Autotuner.autotune_warmup(AutotuneKernelType.DECODE_ATTENTION): - decode_func(copy.copy(infer_state), copy.copy(infer_state1)) + decode_func(_copy_warmup_state(infer_state), _copy_warmup_state(infer_state1)) torch.cuda.synchronize() for para_name in set(vars(infer_state).keys()): if para_name not in pure_para_set: diff --git a/lightllm/common/basemodel/infer_struct.py b/lightllm/common/basemodel/infer_struct.py index 14f7c3f0d9..eed3a757ff 100755 --- a/lightllm/common/basemodel/infer_struct.py +++ b/lightllm/common/basemodel/infer_struct.py @@ -1,5 +1,5 @@ import collections -from typing import TYPE_CHECKING, Tuple, Any, Optional, List +from typing import Tuple, Any, Optional, List import torch import triton @@ -15,9 +15,6 @@ from lightllm.utils.dist_utils import get_global_dp_rank, get_dp_world_size from .attention import BasePrefillAttState, BaseDecodeAttState -if TYPE_CHECKING: - from lightllm.common.linear_state import SSMRunConfig - class InferStateInfo: """ @@ -71,7 +68,6 @@ def __init__(self): # 在microbatch overlap的运行模式下，用于标记当前 microbatch 的 index 序号 # 在一些细节场景下需要有该信息区分一些资源的申请和管理。 self.microbatch_index: int = 0 - self.ssm_run_config: Optional["SSMRunConfig"] = None self.b_logits_row_idx: Optional[torch.Tensor] = None self.input_rows_per_logit: int = 1 diff --git a/lightllm/common/linear_state.py b/lightllm/common/linear_state.py index 1351d6bbb2..bfa913d4fd 100644 --- a/lightllm/common/linear_state.py +++ b/lightllm/common/linear_state.py @@ -29,6 +29,11 @@ def __post_init__(self): raise TypeError("SSM kernel configuration must be an integer") +@dataclass(frozen=True, slots=True) +class LinearStateCommit: + _run_config: SSMRunConfig + + class LinearStateTensors(NamedTuple): conv: torch.Tensor ssm: torch.Tensor @@ -132,6 +137,13 @@ def prepare_mtp_decode(self, reqs, cu_seqlens): self.prepare_decode(reqs, cu_seqlens) return reqs, ssm_indexes, accepted_token_counts + def create_commit(self, run_config: Optional[SSMRunConfig]) -> Optional[LinearStateCommit]: + from lightllm.common.basemodel.triton_kernel.linear_att.compact_ssm import CompactSSMBackend + + if run_config is not None and isinstance(self.ssm_backend, CompactSSMBackend): + return LinearStateCommit(run_config) + return None + def commit_accepted_state( self, b_req_mtp_start_loc, @@ -139,12 +151,12 @@ def commit_accepted_state( b_mtp_index, accepted_index, verify_width, - linear_state_batches: Optional[Sequence[Tuple[int, Optional[SSMRunConfig]]]] = None, + linear_state_batches: Optional[Sequence[Tuple[int, Optional[LinearStateCommit]]]] = None, ): """Commit accepted prefixes without requiring full-state materialization. ``linear_state_batches`` partitions logical requests in the same order as the MTP - start locations. Each configuration belongs to the forward that wrote that + start locations. Each commit belongs to the forward that wrote that partition's scratch; padding rows do not participate in acceptance. """ from lightllm.common.basemodel.triton_kernel.mtp_utils import linear_att_mtp_state_index_update @@ -170,12 +182,12 @@ def commit_accepted_state( return assert sum(size for size, _ in linear_state_batches) == reqs.numel() start = 0 - for size, run_config in linear_state_batches: + for size, commit in linear_state_batches: if size: self.ssm_backend.accept_updates( reqs[start : start + size], self.accepted_state_index, - run_config=run_config, + run_config=commit._run_config if commit is not None else None, ) start += size diff --git a/lightllm/server/router/model_infer/mode_backend/chunked_prefill/impl.py b/lightllm/server/router/model_infer/mode_backend/chunked_prefill/impl.py index ac38e4e975..64174a4fef 100644 --- a/lightllm/server/router/model_infer/mode_backend/chunked_prefill/impl.py +++ b/lightllm/server/router/model_infer/mode_backend/chunked_prefill/impl.py @@ -298,7 +298,7 @@ def decode_mtp( b_req_idx=model_input.b_req_idx, b_req_mtp_start_loc=b_req_mtp_start_loc, b_mtp_index=model_input.b_mtp_index, - linear_state_batches=((req_num, model_output.ssm_run_config),), + linear_state_batches=((req_num, model_output.linear_state_commit),), ) accepted_index_cpu = g_pin_mem_manager.async_copy_from_gpu_tensor( key="accepted_index", diff --git a/lightllm/server/router/model_infer/mode_backend/dp_backend/impl.py b/lightllm/server/router/model_infer/mode_backend/dp_backend/impl.py index ced5d1d64b..f7276e6f0e 100644 --- a/lightllm/server/router/model_infer/mode_backend/dp_backend/impl.py +++ b/lightllm/server/router/model_infer/mode_backend/dp_backend/impl.py @@ -533,7 +533,7 @@ def decode_mtp(self, event_pack: OverlapEventPack, decode_reqs: List[InferReq]): b_req_idx=model_input.b_req_idx, b_req_mtp_start_loc=b_req_mtp_start_loc, b_mtp_index=model_input.b_mtp_index, - linear_state_batches=((req_num, model_output.ssm_run_config),), + linear_state_batches=((req_num, model_output.linear_state_commit),), ) accepted_index_cpu = g_pin_mem_manager.async_copy_from_gpu_tensor( key="accepted_index", @@ -809,8 +809,8 @@ def decode_overlap_mtp(self, event_pack: OverlapEventPack, decode_reqs: List[Inf b_req_mtp_start_loc=b_req_mtp_start_loc, b_mtp_index=b_mtp_index, linear_state_batches=( - (real_request_num0, model_output0.ssm_run_config), - (real_request_num1, model_output1.ssm_run_config), + (real_request_num0, model_output0.linear_state_commit), + (real_request_num1, model_output1.linear_state_commit), ), ) mtp_accept_len0 = mtp_accept_len[:real_request_num0] diff --git a/unit_tests/common/basemodel/test_cuda_graph_autotune.py b/unit_tests/common/basemodel/test_cuda_graph_autotune.py index 52b1c1cf2d..b5b50beb20 100644 --- a/unit_tests/common/basemodel/test_cuda_graph_autotune.py +++ b/unit_tests/common/basemodel/test_cuda_graph_autotune.py @@ -1,5 +1,6 @@ from contextlib import contextmanager from types import SimpleNamespace +import weakref import pytest import torch @@ -128,3 +129,59 @@ def decode_func(*infer_states): assert not Autotuner.is_autotune_warmup() assert capture_env.capture_count == 0 assert graph.graph == {} + + +@pytest.mark.parametrize("overlap", [False, True]) +@pytest.mark.parametrize("kda", [False, True]) +def test_linear_config_is_local_to_capture_state(capture_env, overlap, kda): + from lightllm.common.basemodel.attention.linear.gdn import GDNLinearAttDecodeAttState + from lightllm.common.basemodel.attention.linear.kda import KDALinearAttDecodeAttState + from lightllm.common.basemodel.attention.triton.fp import TritonDecodeAttState + from lightllm.common.basemodel.basemodel import TpPartBaseModel + from lightllm.common.basemodel.batch_objs import PostLayerOutput + from lightllm.common.basemodel.hidden_collector import NoopHiddenCollector + from lightllm.common.basemodel.infer_struct import InferStateInfo + from lightllm.common.basemodel.triton_kernel.linear_att.compact_ssm import CompactSSMBackend + from lightllm.common.linear_state import LinearStateManager, SSMRunConfig + + capture_env.args.enable_decode_microbatch_overlap = overlap + graph = CudaGraph(1, 2, 1, max_batch_size=2) + manager = object.__new__(LinearStateManager) + manager.ssm_backend = object.__new__(CompactSSMBackend) + att_class = KDALinearAttDecodeAttState if kda else GDNLinearAttDecodeAttState + states = [] + configs = [SSMRunConfig(8, 1), SSMRunConfig(32, 2)] + for _ in range(2 if overlap else 1): + state = InferStateInfo() + state.input_ids = torch.ones(2, dtype=torch.int64) + state.req_manager = SimpleNamespace(linear_state_manager=manager) + state.hidden_collector = NoopHiddenCollector() + state.decode_att_state = TritonDecodeAttState(_infer_state_weak_ref=weakref.ref(state)) + state.decode_att_state1 = att_class(_infer_state_weak_ref=weakref.ref(state)) + states.append(state) + + def decode_func(*infer_states): + outputs = [] + for index, state in enumerate(infer_states): + att = state.decode_att_state1 + assert att.infer_state is state + assert att.run_config is None + assert (att is states[index].decode_att_state1) == capture_env.capturing + att.run_config = configs[index] if capture_env.capturing else SSMRunConfig(128, 4) + outputs.append(TpPartBaseModel._create_model_output(None, PostLayerOutput(torch.ones(2, 1)), state)) + return tuple(outputs) if overlap else outputs[0] + + result = graph.capture_decode(decode_func, *states) + outputs = result if overlap else (result,) + runtime_states = [cuda_graph_module._copy_warmup_state(state) for state in states] + for state in runtime_states: + state.decode_att_state1.run_config = None + replay = graph.replay(*runtime_states) + assert replay == result + for index, (state, output) in enumerate(zip(states, outputs)): + assert state.decode_att_state1.run_config is configs[index] + state.decode_att_state1.run_config = SSMRunConfig(64, 4) + assert output.linear_state_commit == manager.create_commit(configs[index]) + unpadded = TpPartBaseModel._create_unpad_decode_model_output(None, output, 1) + assert unpadded.logits.shape[0] == 1 + assert unpadded.linear_state_commit is output.linear_state_commit diff --git a/unit_tests/common/basemodel/triton_kernel/linear_att/test_ssm_autotune.py b/unit_tests/common/basemodel/triton_kernel/linear_att/test_ssm_autotune.py index e824ad77c1..123c2cc910 100644 --- a/unit_tests/common/basemodel/triton_kernel/linear_att/test_ssm_autotune.py +++ b/unit_tests/common/basemodel/triton_kernel/linear_att/test_ssm_autotune.py @@ -319,6 +319,10 @@ def choose(ssm_backend, mode, q, *args): @pytest.mark.parametrize("dynamic", [False, True]) def test_two_verify_graphs_accept_with_their_own_configs(monkeypatch, mode, projection, dtype, kda, dynamic): from types import SimpleNamespace + import weakref + from lightllm.common.basemodel.infer_struct import InferStateInfo + from lightllm.common.basemodel.attention.linear.gdn import GDNLinearAttDecodeAttState + from lightllm.common.basemodel.attention.linear.kda import KDALinearAttDecodeAttState from lightllm.common.basemodel.basemodel import TpPartBaseModel from lightllm.common.basemodel.batch_objs import PostLayerOutput, ModelMtpOutputCollector from lightllm.common.req_manager.linear_att import ReqManagerForMamba @@ -366,17 +370,18 @@ def forward(target): graph = torch.cuda.CUDAGraph() with torch.cuda.graph(graph): outputs, config = forward(ssm_backend) - state = SimpleNamespace( - hidden_collector=SimpleNamespace(finish_output=lambda **kw: ModelMtpOutputCollector()), - prompt_logics=None, - ssm_run_config=config, - input_rows_per_logit=1, - ) + state = InferStateInfo() + state.req_manager = manager + state.hidden_collector = SimpleNamespace(finish_output=lambda **kw: ModelMtpOutputCollector()) + att_class = KDALinearAttDecodeAttState if kda else GDNLinearAttDecodeAttState + state.decode_att_state1 = att_class(_infer_state_weak_ref=weakref.ref(state), run_config=config) output = TpPartBaseModel._create_model_output( None, PostLayerOutput(logits=outputs[0].view(count * 4, -1)), state ) + commit = output.linear_state_commit + assert (commit is not None) == (mode == "compact") output = TpPartBaseModel._create_unpad_decode_model_output(None, output, len(ids) * 4) - assert output.ssm_run_config is config + assert output.linear_state_commit is commit buckets.append((graph, outputs, output, ref, args, reqs, cu, ids)) # Published contexts and captured launches must survive later changes to # the tuner's dictionaries, including when two graphs replay alternately. @@ -406,7 +411,7 @@ def forward(target): mtp, flags, 4, - linear_state_batches=tuple((len(b[-1]), b[2].ssm_run_config) for b in buckets), + linear_state_batches=tuple((len(b[-1]), b[2].linear_state_commit) for b in buckets), ) for _, _, _, ref, _, reqs, _, ids in buckets: ref.accept_updates(reqs[:-1], manager.linear_state_manager.accepted_state_index) diff --git a/unit_tests/common/test_linear_state_contract.py b/unit_tests/common/test_linear_state_contract.py index 22ac52b792..32b1f98c33 100644 --- a/unit_tests/common/test_linear_state_contract.py +++ b/unit_tests/common/test_linear_state_contract.py @@ -118,3 +118,33 @@ def test_ssm_backend_publishes_detached_run_config(mode): assert config.num_stages == 3 with pytest.raises(FrozenInstanceError): config.BV = 128 + + +def test_compact_commits_each_unpadded_partition_with_its_verify_config(monkeypatch): + from lightllm.common.basemodel.triton_kernel import mtp_utils + from lightllm.common.basemodel.triton_kernel.linear_att.compact_ssm import CompactSSMBackend + + owner = object.__new__(LinearStateManager) + owner.ssm_backend = object.__new__(CompactSSMBackend) + owner.accepted_state_index = torch.tensor([0, 1, 2, 0], dtype=torch.int32) + configs = [SSMRunConfig(8, 1), SSMRunConfig(32, 2)] + commits = [owner.create_commit(config) for config in configs] + owner.ssm_backend.run_config = {"BV": 128, "num_warps": 4} + calls = [] + monkeypatch.setattr(mtp_utils, "linear_att_mtp_state_index_update", lambda **kwargs: None) + monkeypatch.setattr( + owner.ssm_backend, + "accept_updates", + lambda reqs, accepted, run_config: calls.append((reqs.tolist(), accepted, run_config)), + ) + owner.commit_accepted_state( + torch.tensor([0, 4, 8]), + torch.tensor([2] * 4 + [0] * 4 + [1] * 4), + torch.arange(4).repeat(3), + torch.ones(12, dtype=torch.int32), + 4, + linear_state_batches=((1, commits[0]), (0, None), (2, commits[1])), + ) + assert [reqs for reqs, _, _ in calls] == [[2], [0, 1]] + assert all(accepted is owner.accepted_state_index for _, accepted, _ in calls) + assert [config for _, _, config in calls] == configs diff --git a/unit_tests/models/test_post_layer_output.py b/unit_tests/models/test_post_layer_output.py index 9e56e47055..4febaca2e1 100644 --- a/unit_tests/models/test_post_layer_output.py +++ b/unit_tests/models/test_post_layer_output.py @@ -291,7 +291,11 @@ def test_model_output_and_unpadding_preserve_candidate_mapping() -> None: logits = torch.randn(3, 2) ids = torch.tensor([[10, 20], [30, 40], [50, 60]]) state = SimpleNamespace( - hidden_collector=NoopHiddenCollector(), prompt_logics=None, ssm_run_config=None, input_rows_per_logit=1 + hidden_collector=NoopHiddenCollector(), + prompt_logics=None, + decode_att_state=None, + decode_att_state1=None, + input_rows_per_logit=1, ) output = model._create_model_output(PostLayerOutput(logits, ids), state) unpadded = model._create_unpad_decode_model_output(output, 2) From ac4a82c52713da4d7c412b8ddb74d07d2a660133 Mon Sep 17 00:00:00 2001 From: wangzaijun Date: Thu, 22 Oct 2026 06:20:22 +0000 Subject: [PATCH 34/40] refactor: split accepted-tail draft projection out of ReplaySSM PR --- lightllm/common/basemodel/basemodel.py | 63 +------ lightllm/common/basemodel/batch_objs.py | 7 - lightllm/common/basemodel/cuda_graph.py | 6 +- lightllm/common/basemodel/infer_struct.py | 2 - .../llama/layer_infer/post_layer_infer.py | 3 - .../proposers/eagle_with_att.py | 14 +- .../common/basemodel/test_draft_tail_graph.py | 167 ------------------ .../common/basemodel/test_model_output.py | 35 ---- unit_tests/models/test_post_layer_output.py | 3 +- .../mtp_speculative/test_eagle_with_att.py | 48 +---- 10 files changed, 14 insertions(+), 334 deletions(-) delete mode 100644 unit_tests/common/basemodel/test_draft_tail_graph.py diff --git a/lightllm/common/basemodel/basemodel.py b/lightllm/common/basemodel/basemodel.py index 1488937daf..2c42547659 100755 --- a/lightllm/common/basemodel/basemodel.py +++ b/lightllm/common/basemodel/basemodel.py @@ -360,21 +360,6 @@ def _align_decode_batch_size(self, batch_size: int) -> int: alignment = math.lcm(alignment, self.tp_world_size_) return triton.cdiv(batch_size, alignment) * alignment - def should_project_draft_tail_logits(self): - """Project only accepted-tail rows in the first draft pass.""" - return ( - self.is_mtp_draft_model - and self.args.ssm_state_mode == "replay" - and self.args.ssm_replay_config.get("compute_mode", "inline") == "parallel" - and self.args.mtp_mode == "eagle_with_att" - and self.args.mtp_step > 0 - and not self.args.mtp_dynamic_verify - and not self.args.enable_decode_microbatch_overlap - and not self.args.enable_tpsp_mix_mode - and self.tp_world_size_ == 1 - and self.config.get("model_type") in ("qwen3_5", "qwen3_5_text") - ) - def _init_cudagraph(self): # When graph covers the configured request length, it must also cover MTP's internal token margin. if self.args.mtp_mode is not None and self.graph_max_len_in_batch >= self.args.max_req_total_len: @@ -395,21 +380,7 @@ def _init_cudagraph(self): capture_infer_cost=self.args.mtp_dynamic_verify, ) ) - self.draft_tail_graph = None if self.graph is not None: - if self.should_project_draft_tail_logits(): - width = self.args.mtp_step + 1 - self.draft_tail_graph = CudaGraph( - batch_step_size_before_split=width, - split_batch_size=self.args.graph_split_batch_size * width, - batch_step_size_after_split=self.args.graph_grow_step_size * width, - max_batch_size=self.graph_max_batch_size * width, - max_len_in_batch=self.graph_max_len_in_batch, - tp_world_size=self.tp_world_size_, - ) - # Draft passes execute sequentially, so both layouts can share the pool. - self.draft_tail_graph.mempool = self.graph.mempool - self.draft_tail_graph.warmup(self, input_rows_per_logit=width) if get_env_start_args().enable_decode_microbatch_overlap: self.graph.warmup_overlap(self) else: @@ -468,16 +439,6 @@ def _create_inferstate(self, model_input: ModelInput, microbatch_index: int = 0) infer_state.att_state_role = self.att_state_role infer_state.return_all_prompt_logics = self.return_all_prompt_logics infer_state.batch_size = model_input.batch_size - if model_input.b_logits_row_idx is not None: - assert self.should_project_draft_tail_logits() and not model_input.is_prefill - width = self.args.mtp_step + 1 - assert model_input.batch_size % width == 0 - count = model_input.batch_size // width - rows = model_input.b_logits_row_idx - assert rows.ndim == 1 and rows.numel() <= count - rows = F.pad(rows, (0, count - rows.numel())) - infer_state.b_logits_row_idx = rows - infer_state.input_rows_per_logit = width infer_state.total_token_num = model_input.total_token_num infer_state.max_q_seq_len = model_input.max_q_seq_len infer_state.max_kv_seq_len = model_input.max_kv_seq_len @@ -631,21 +592,17 @@ def _create_model_output(self, post_output: PostLayerOutput, infer_state: InferS prompt_logics=infer_state.prompt_logics, logits_token_ids=post_output.logits_token_ids, linear_state_commit=linear_state_commit, - input_rows_per_logit=infer_state.input_rows_per_logit, ) return output def _create_unpad_decode_model_output(self, model_output: ModelOutput, origin_batch_size: int): - input_rows_per_logit = model_output.input_rows_per_logit - padded_batch_size = model_output.logits.shape[0] * input_rows_per_logit - assert origin_batch_size % input_rows_per_logit == 0 - logits_size = origin_batch_size // input_rows_per_logit + padded_batch_size = model_output.logits.shape[0] if padded_batch_size == origin_batch_size: return model_output new_model_output = copy.copy(model_output) - new_model_output.logits = new_model_output.logits[0:logits_size] + new_model_output.logits = new_model_output.logits[0:origin_batch_size] if new_model_output.logits_token_ids is not None: - new_model_output.logits_token_ids = new_model_output.logits_token_ids[:logits_size] + new_model_output.logits_token_ids = new_model_output.logits_token_ids[:origin_batch_size] new_model_output.mtp_collector = model_output.mtp_collector.unpad_decode( padded_batch_size=padded_batch_size, origin_batch_size=origin_batch_size, @@ -739,24 +696,20 @@ def _decode( ) origin_batch_size = model_input.batch_size - graph = self.draft_tail_graph if model_input.b_logits_row_idx is not None else self.graph # 空 DP rank 也需要完整的 dummy verify group；TP/SP 切分不能破坏 MTP 分组。 infer_batch_size = self._align_decode_batch_size(max(1, origin_batch_size)) - if model_input.b_logits_row_idx is not None: - width = self.args.mtp_step + 1 - infer_batch_size = triton.cdiv(infer_batch_size, width) * width # CUDA Graph 可能继续向上对齐 batch size，并因此加入 seq_len=2 的 # dummy request。先用最终可能出现的 KV 长度判断 graph，再统一 padding 一次。 infer_max_kv_seq_len = max(2, model_input.max_kv_seq_len) - use_cuda_graph = graph is not None and graph.can_run( + use_cuda_graph = self.graph is not None and self.graph.can_run( batch_size=infer_batch_size, max_len_in_batch=infer_max_kv_seq_len, ) need_capture = False if use_cuda_graph: - infer_batch_size = graph.find_closest_graph_batch_size(batch_size=infer_batch_size) - need_capture = graph.need_capture(infer_batch_size) + infer_batch_size = self.graph.find_closest_graph_batch_size(batch_size=infer_batch_size) + need_capture = self.graph.need_capture(infer_batch_size) model_input = self._create_padded_decode_model_input(model_input=model_input, new_batch_size=infer_batch_size) infer_state = self._create_inferstate(model_input) @@ -769,9 +722,9 @@ def _decode( if use_cuda_graph: if need_capture: - model_output: ModelOutput = graph.capture_decode(self._token_forward, infer_state) + model_output: ModelOutput = self.graph.capture_decode(self._token_forward, infer_state) else: - model_output: ModelOutput = graph.replay(infer_state) + model_output: ModelOutput = self.graph.replay(infer_state) else: model_output = self._token_forward(infer_state) diff --git a/lightllm/common/basemodel/batch_objs.py b/lightllm/common/basemodel/batch_objs.py index 3b1a27ebc8..b2e2b54752 100644 --- a/lightllm/common/basemodel/batch_objs.py +++ b/lightllm/common/basemodel/batch_objs.py @@ -55,13 +55,8 @@ class ModelInput: # 的 draft 模型的输入 mtp_draft_input_hiddens: Optional[torch.Tensor] = None - # Optional accepted-tail rows for the first EAGLE draft projection. - b_logits_row_idx: Optional[torch.Tensor] = None - def to_cuda(self): self.check_input() - if self.b_logits_row_idx is not None: - self.b_logits_row_idx = self.b_logits_row_idx.cuda(non_blocking=True) # Prefill 和 decode 都必须提供的公共张量。 self.b_req_idx = self.b_req_idx.cuda(non_blocking=True) @@ -212,8 +207,6 @@ class ModelOutput: # Detached state-commit metadata for this forward/graph; preserved when unpadding. linear_state_commit: Optional["LinearStateCommit"] = None - # Number of physical decode rows represented by each selected logits row. - input_rows_per_logit: int = 1 def __post_init__(self) -> None: if self.mtp_collector is None: diff --git a/lightllm/common/basemodel/cuda_graph.py b/lightllm/common/basemodel/cuda_graph.py index 21cf83e5f4..39de136d40 100644 --- a/lightllm/common/basemodel/cuda_graph.py +++ b/lightllm/common/basemodel/cuda_graph.py @@ -264,7 +264,7 @@ def replay(self, infer_state, infer_state1=None): return self._replay(infer_state) @torch.no_grad() - def warmup(self, model, input_rows_per_logit=1): + def warmup(self, model): logger.info("Begin capture cudagraph, use the --disable_cudagraph to disable it.") # for typing easy from .basemodel import TpPartBaseModel @@ -300,10 +300,6 @@ def warmup(self, model, input_rows_per_logit=1): multimodal_params=[{"images": [], "audios": []} for _ in range(batch_size)], **model._gen_special_model_input(batch_size), ) - if input_rows_per_logit > 1: - model_input.b_logits_row_idx = torch.arange( - input_rows_per_logit - 1, batch_size, input_rows_per_logit, dtype=torch.int64, device="cuda" - ) model_output: ModelOutput = model.forward(model_input) del model_output del input_ids diff --git a/lightllm/common/basemodel/infer_struct.py b/lightllm/common/basemodel/infer_struct.py index eed3a757ff..9f364da40b 100755 --- a/lightllm/common/basemodel/infer_struct.py +++ b/lightllm/common/basemodel/infer_struct.py @@ -68,8 +68,6 @@ def __init__(self): # 在microbatch overlap的运行模式下，用于标记当前 microbatch 的 index 序号 # 在一些细节场景下需要有该信息区分一些资源的申请和管理。 self.microbatch_index: int = 0 - self.b_logits_row_idx: Optional[torch.Tensor] = None - self.input_rows_per_logit: int = 1 # 当前 forward 独占的 hidden state 收集器。普通推理使用短生命周期实例； # Prefill CUDA Graph 使用随 graph infer state 长期保存的实例。 diff --git a/lightllm/models/llama/layer_infer/post_layer_infer.py b/lightllm/models/llama/layer_infer/post_layer_infer.py index 93fe21974a..0e0ed924d2 100644 --- a/lightllm/models/llama/layer_infer/post_layer_infer.py +++ b/lightllm/models/llama/layer_infer/post_layer_infer.py @@ -242,9 +242,6 @@ def _draft_lm_head_and_gather( ) -> PostLayerOutput: """按 draft 配置返回候选 logits，未开启时回退到完整词表。""" - if infer_state.b_logits_row_idx is not None: - hidden = hidden.index_select(0, infer_state.b_logits_row_idx) - token_num = hidden.shape[0] top_k = get_env_start_args().draft_vocab_topk_sampling if top_k is None: return self._lm_head_and_gather(hidden, token_num, layer_weight, infer_state) diff --git a/lightllm/server/router/model_infer/mtp_speculative/proposers/eagle_with_att.py b/lightllm/server/router/model_infer/mtp_speculative/proposers/eagle_with_att.py index 4f84b7e8c7..721af1c27d 100644 --- a/lightllm/server/router/model_infer/mtp_speculative/proposers/eagle_with_att.py +++ b/lightllm/server/router/model_infer/mtp_speculative/proposers/eagle_with_att.py @@ -73,24 +73,14 @@ def propose_next( verify_draft_input = copy.copy(target_model_input) verify_draft_input.input_ids = target_next_token_ids verify_draft_input.mtp_draft_input_hiddens = target_model_output.mtp_collector.spec_hidden - if draft_model.should_project_draft_tail_logits(): - verify_draft_input.b_logits_row_idx = accepted_tail_rows extend_output = draft_model.forward(verify_draft_input) # 只在 req_num 行 logits 上进行 argmax，避免为未接受的 verify 行执行 # vocabulary reduction。第一列 proposal 来自每个请求的 accepted tail。 accepted_tail_output = ModelOutput( - logits=( - extend_output.logits[:req_num] - if extend_output.input_rows_per_logit > 1 - else extend_output.logits.index_select(0, accepted_tail_rows) - ), + logits=extend_output.logits.index_select(0, accepted_tail_rows), logits_token_ids=( - ( - extend_output.logits_token_ids[:req_num] - if extend_output.input_rows_per_logit > 1 - else extend_output.logits_token_ids.index_select(0, accepted_tail_rows) - ) + extend_output.logits_token_ids.index_select(0, accepted_tail_rows) if extend_output.logits_token_ids is not None else None ), diff --git a/unit_tests/common/basemodel/test_draft_tail_graph.py b/unit_tests/common/basemodel/test_draft_tail_graph.py deleted file mode 100644 index 57ce202959..0000000000 --- a/unit_tests/common/basemodel/test_draft_tail_graph.py +++ /dev/null @@ -1,167 +0,0 @@ -from types import SimpleNamespace - -import pytest -import torch - -from lightllm.common.basemodel import basemodel, cuda_graph -from lightllm.common.basemodel.basemodel import TpPartBaseModel -from lightllm.common.basemodel.batch_objs import ModelInput, ModelMtpOutputCollector, ModelOutput -from lightllm.common.basemodel.hidden_collector import NoopHiddenCollector -from lightllm.common.basemodel.infer_struct import InferStateInfo -from lightllm.server.core.objs.start_args_type import StartArgs - - -@pytest.fixture -def draft_model(monkeypatch): - args = StartArgs( - ssm_state_mode="replay", - ssm_replay_config={"compute_mode": "parallel"}, - mtp_mode="eagle_with_att", - mtp_step=4, - running_max_req_size=264, - graph_max_batch_size=256, - max_req_total_len=128, - ) - monkeypatch.setattr(basemodel, "get_env_start_args", lambda: args) - monkeypatch.setattr(cuda_graph, "get_env_start_args", lambda: args) - monkeypatch.setattr(basemodel.dist_group_manager, "get_group", lambda _: None) - monkeypatch.setattr(InferStateInfo, "init_some_extra_state", lambda *args: None) - warmups = [] - monkeypatch.setattr(cuda_graph.CudaGraph, "warmup", lambda graph, model, **kw: warmups.append((graph, kw))) - - # Execute CPU row selection while retaining the real graph lookup, input - # copying, padding and output unpadding paths. - def capture(graph, forward, state): - output = forward(state) - - def replay(): - current = forward(state) - output.logits.copy_(current.logits) - output.mtp_collector.spec_hidden.copy_(current.mtp_collector.spec_hidden) - - graph.graph[state.batch_size] = (SimpleNamespace(replay=replay), state, output) - return output - - monkeypatch.setattr(cuda_graph.CudaGraph, "capture_decode", capture) - model = TpPartBaseModel.__new__(TpPartBaseModel) - model.args = args - model.is_mtp_draft_model = True - model.tp_world_size_ = 1 - model.config = {"model_type": "qwen3_5_text"} - model.disable_cudagraph = False - model.graph_max_batch_size = args.graph_max_batch_size - model.graph_max_len_in_batch = 128 - model.max_seq_length = 128 - model.mtp_manager = SimpleNamespace( - get_decode_tokens_per_request=lambda _: 1, get_decode_batch_alignment=lambda _: 1 - ) - model.infer_state_class = InferStateInfo - model.hidden_collector_prototype = NoopHiddenCollector() - model.return_all_prompt_logics = False - model.req_manager = SimpleNamespace(HOLD_REQUEST_ID=264) - model.mem_manager = None - model._select_mem_indexes = lambda _: None - model.decode_att_backend = SimpleNamespace( - create_att_decode_state=lambda **kw: SimpleNamespace( - init_state=lambda: None, copy_for_decode_cuda_graph=lambda _: None, check_no_infer_state=lambda: None - ) - ) - model.decode_att_backend1 = None - - def forward(state): - hidden = state.input_ids[:, None].float() - logits = hidden if state.b_logits_row_idx is None else hidden.index_select(0, state.b_logits_row_idx) - return ModelOutput( - logits=logits, - input_rows_per_logit=state.input_rows_per_logit, - mtp_collector=ModelMtpOutputCollector(spec_hidden=hidden), - ) - - model._token_forward = forward - model._init_cudagraph() - assert warmups == [(model.draft_tail_graph, {"input_rows_per_logit": 5}), (model.graph, {})] - return model - - -def make_input(rows, accepted=None): - return ModelInput( - batch_size=rows, - total_token_num=rows * 10, - max_q_seq_len=1, - max_kv_seq_len=10, - input_ids=torch.arange(rows), - b_req_idx=torch.arange(rows, dtype=torch.int32) // (5 if accepted is not None else 1), - b_mtp_index=torch.zeros(rows, dtype=torch.int32), - b_seq_len=torch.full((rows,), 10, dtype=torch.int32), - b_position_delta=torch.zeros(rows, dtype=torch.int32), - b_shared_seq_len=torch.zeros(rows, dtype=torch.int32), - b_shared_radix_node_id=torch.zeros(rows, dtype=torch.int64), - multimodal_params=[{} for _ in range(rows)], - b_logits_row_idx=accepted, - ) - - -@pytest.mark.parametrize("eager", [False, True]) -@pytest.mark.parametrize("requests", [0, 3, 32, 51, 52, 53, 264]) -def test_first_draft_selects_accepted_tail_across_padding(draft_model, eager, requests): - model = draft_model - if eager: - model.graph = model.draft_tail_graph = None - for step in range(2): - accepted = torch.arange(requests) * 5 + (torch.arange(requests) + step) % 5 - model_input = make_input(requests * 5, accepted) - output = model._decode(model_input) - torch.testing.assert_close(output.logits[:, 0], accepted.float()) - torch.testing.assert_close(output.mtp_collector.spec_hidden[:, 0], model_input.input_ids.float()) - assert output.input_rows_per_logit == 5 - assert model_input.batch_size == requests * 5 - assert model_input.b_logits_row_idx is accepted - if not eager and requests == 52: - assert 320 in model.draft_tail_graph.graph - assert not model.graph.graph - - -@pytest.mark.parametrize("rows", [0, 160, 256, 260, 263, 264]) -def test_later_draft_keeps_one_output_per_request(draft_model, rows): - model = draft_model - # 160 physical rows occur in both layouts; their captured outputs must not mix. - model._decode(make_input(160, torch.arange(32) * 5 + 1)) - for step in range(2): - model_input = make_input(rows) - model_input.input_ids += step - output = model._decode(model_input) - assert output.input_rows_per_logit == 1 - torch.testing.assert_close(output.logits[:, 0], model_input.input_ids.float()) - assert output.mtp_collector.spec_hidden.shape == (rows, 1) - if rows == 160: - assert 160 in model.graph.graph and 160 in model.draft_tail_graph.graph - assert model.graph.graph[160][2].logits.shape == (160, 1) - assert model.draft_tail_graph.graph[160][2].logits.shape == (32, 1) - elif rows > 256: - assert not model.graph.graph - - -def test_tail_graph_warmup_marks_selected_rows(monkeypatch): - args = StartArgs() - monkeypatch.setattr(cuda_graph, "get_env_start_args", lambda: args) - - def on_cpu(factory): - def allocate(*args, **kwargs): - kwargs["device"] = "cpu" - return factory(*args, **kwargs) - - return allocate - - for name in ("tensor", "full", "zeros", "arange"): - monkeypatch.setattr(torch, name, on_cpu(getattr(torch, name))) - seen = [] - model = SimpleNamespace( - req_manager=SimpleNamespace(HOLD_REQUEST_ID=8, free_all=lambda: None), - mem_manager=SimpleNamespace(free_all=lambda: None), - _gen_special_model_input=lambda _: {}, - forward=lambda model_input: seen.append(model_input), - ) - graph = cuda_graph.CudaGraph(5, 5, 5, max_batch_size=10) - graph.warmup(model, input_rows_per_logit=5) - assert [x.batch_size for x in seen] == [10, 5] - assert [x.b_logits_row_idx.tolist() for x in seen] == [[4, 9], [4]] diff --git a/unit_tests/common/basemodel/test_model_output.py b/unit_tests/common/basemodel/test_model_output.py index 7a9045a344..123f603ff2 100644 --- a/unit_tests/common/basemodel/test_model_output.py +++ b/unit_tests/common/basemodel/test_model_output.py @@ -326,41 +326,6 @@ def forward(state): torch.testing.assert_close(torch.stack((model_input.b_req_idx, model_input.b_seq_len), dim=1), expected) -@pytest.mark.parametrize("origin", [0, 5, 10, 15]) -def test_selected_draft_logits_unpad_preserves_physical_hidden_rows(origin): - model = TpPartBaseModel.__new__(TpPartBaseModel) - output = ModelOutput( - logits=torch.arange(12).view(3, 4), - logits_token_ids=torch.arange(12).view(3, 4), - mtp_collector=ModelMtpOutputCollector(spec_hidden=torch.arange(45).view(15, 3)), - input_rows_per_logit=5, - ) - unpadded = model._create_unpad_decode_model_output(output, origin) - assert unpadded.logits.shape == (origin // 5, 4) - assert unpadded.logits_token_ids.shape == (origin // 5, 4) - assert unpadded.mtp_collector.spec_hidden.shape == (origin, 3) - assert output.mtp_collector.spec_hidden.shape == (15, 3) - assert unpadded.input_rows_per_logit == 5 - - -@pytest.mark.parametrize("disabled", [None, "dynamic", "overlap", "tp", "projection", "model", "target"]) -def test_draft_tail_scope(disabled): - from lightllm.server.core.objs.start_args_type import StartArgs - - model = TpPartBaseModel.__new__(TpPartBaseModel) - model.args = StartArgs( - ssm_state_mode="replay", ssm_replay_config={"compute_mode": "parallel"}, mtp_mode="eagle_with_att", mtp_step=4 - ) - model.is_mtp_draft_model = disabled != "target" - model.tp_world_size_ = 2 if disabled == "tp" else 1 - model.config = {"model_type": "qwen3_5_text" if disabled != "model" else "qwen3_next"} - model.args.mtp_dynamic_verify = disabled == "dynamic" - model.args.enable_decode_microbatch_overlap = disabled == "overlap" - if disabled == "projection": - model.args.ssm_replay_config = {"compute_mode": "inline"} - assert model.should_project_draft_tail_logits() == (disabled is None) - - def _create_empty_decode_input(): return ModelInput( batch_size=0, diff --git a/unit_tests/models/test_post_layer_output.py b/unit_tests/models/test_post_layer_output.py index 4febaca2e1..a2af828623 100644 --- a/unit_tests/models/test_post_layer_output.py +++ b/unit_tests/models/test_post_layer_output.py @@ -60,7 +60,7 @@ def candidates(*args): args = SimpleNamespace(target_vocab_topk_sampling=None, draft_vocab_topk_sampling=None) setattr(args, "draft_vocab_topk_sampling" if draft else "target_vocab_topk_sampling", 16 if enabled else None) monkeypatch.setattr(llama_post, "get_env_start_args", lambda: args) - state = SimpleNamespace(is_draft_model=draft, dist_group=None, b_logits_row_idx=None) + state = SimpleNamespace(is_draft_model=draft, dist_group=None) weight = SimpleNamespace(lm_head_weight_=Weight()) method = head._draft_lm_head_and_gather if draft else head._target_lm_head_and_gather output = method(torch.ones(3, 2), 3, weight, state) @@ -295,7 +295,6 @@ def test_model_output_and_unpadding_preserve_candidate_mapping() -> None: prompt_logics=None, decode_att_state=None, decode_att_state1=None, - input_rows_per_logit=1, ) output = model._create_model_output(PostLayerOutput(logits, ids), state) unpadded = model._create_unpad_decode_model_output(output, 2) diff --git a/unit_tests/server/router/model_infer/mtp_speculative/test_eagle_with_att.py b/unit_tests/server/router/model_infer/mtp_speculative/test_eagle_with_att.py index 06da6896b4..365acb30b2 100644 --- a/unit_tests/server/router/model_infer/mtp_speculative/test_eagle_with_att.py +++ b/unit_tests/server/router/model_infer/mtp_speculative/test_eagle_with_att.py @@ -23,7 +23,6 @@ def forward(model_input): draft_model = SimpleNamespace( forward=forward, - should_project_draft_tail_logits=lambda: False, map_draft_vocab_to_main_vocab=lambda token_ids: token_ids + 100, ) proposer = Eagle3Proposer( @@ -103,9 +102,7 @@ def forward(model_input): return ModelOutput(logits=torch.empty((0,), device=device)) proposer = EagleWithAttProposer( - backend=SimpleNamespace( - draft_models=[SimpleNamespace(forward=forward, should_project_draft_tail_logits=lambda: False)] - ), + backend=SimpleNamespace(draft_models=[SimpleNamespace(forward=forward)]), enable_dynmaic_mtp=False, ) proposer.fill_draft_model_kv_state( @@ -161,7 +158,7 @@ def forward(model_input): ) backend = SimpleNamespace( - draft_models=[SimpleNamespace(forward=forward, should_project_draft_tail_logits=lambda: False)], + draft_models=[SimpleNamespace(forward=forward)], _gen_argmax_token_ids=lambda output: output.logits[:, 0].long(), _gen_argmax_token_ids_and_prob=lambda output: ( output.logits[:, 0].long(), @@ -214,44 +211,3 @@ def forward(model_input): assert draft_calls[1]["total_token_num"] == 46 assert target_input.input_ids is original_input_ids assert target_input.mtp_draft_input_hiddens is None - - -@pytest.mark.parametrize("capacity", [2, 10, 264]) -def test_selected_tail_logits_keep_complete_draft_kv_input(capacity): - hidden = torch.arange(20).view(10, 2).float() - target_input = SimpleNamespace( - is_prefill=False, - batch_size=10, - input_ids=None, - b_position_delta=torch.zeros(10, dtype=torch.int32), - mtp_draft_input_hiddens=None, - b_logits_row_idx=None, - ) - - def forward(model_input): - assert model_input.batch_size == 10 - assert model_input.mtp_draft_input_hiddens is hidden - assert model_input.input_ids.numel() == 10 - torch.testing.assert_close(model_input.b_logits_row_idx, torch.tensor([1, 8])) - return ModelOutput(logits=torch.tensor([[7.0], [9.0]]), input_rows_per_logit=5) - - draft_model = SimpleNamespace( - forward=forward, - should_project_draft_tail_logits=lambda: True, - args=SimpleNamespace(running_max_req_size=capacity), - ) - proposer = EagleWithAttProposer( - backend=SimpleNamespace(draft_models=[draft_model], _gen_argmax_token_ids=lambda out: out.logits[:, 0].long()), - enable_dynmaic_mtp=False, - ) - output = proposer.propose_next( - target_input, - ModelOutput(logits=torch.empty(10, 1), mtp_collector=ModelMtpOutputCollector(spec_hidden=hidden)), - torch.arange(10), - torch.tensor([0, 5]), - 1, - torch.tensor([2, 4]), - ) - torch.testing.assert_close(output.token_ids, torch.tensor([[7], [9]])) - assert target_input.b_logits_row_idx is None - assert target_input.input_ids is None From 355357a95c964d5e0182aada71aea3251146ea85 Mon Sep 17 00:00:00 2001 From: wangzaijun Date: Thu, 22 Oct 2026 06:34:55 +0000 Subject: [PATCH 35/40] style: restore original infer_struct imports --- lightllm/common/basemodel/infer_struct.py | 4 +--- 1 file changed, 1 insertion(+), 3 deletions(-) diff --git a/lightllm/common/basemodel/infer_struct.py b/lightllm/common/basemodel/infer_struct.py index 9f364da40b..83b312f991 100755 --- a/lightllm/common/basemodel/infer_struct.py +++ b/lightllm/common/basemodel/infer_struct.py @@ -1,12 +1,10 @@ -import collections -from typing import Tuple, Any, Optional, List - import torch import triton import collections from lightllm.common.attention_state.manager import AttentionStateManager from lightllm.common.req_manager import ReqManager from lightllm.distributed import CustomProcessGroup +from typing import Tuple, Any, Optional, List from .triton_kernel.gen_prefill_params import gen_prefill_params from .triton_kernel.gen_decode_params import gen_decode_params from .triton_kernel.multimodal_emb import mark_multimodal_obj From 69f65be289e680163ebfd1e81651d6f46757e98e Mon Sep 17 00:00:00 2001 From: wangzaijun Date: Thu, 22 Oct 2026 07:38:29 +0000 Subject: [PATCH 36/40] refactor: encapsulate hybrid checkpoint and PD state interfaces --- .../bounded_att/gdn/create_gdn_state.py | 7 + .../bounded_att/gdn/gdn_state.py | 96 ++-- .../bounded_att/kda/create_kda_state.py | 7 + .../bounded_att/kda/kda_state.py | 96 ++-- .../common/basemodel/attention/linear/gdn.py | 21 +- .../common/basemodel/attention/linear/kda.py | 13 +- lightllm/common/basemodel/basemodel.py | 2 - .../triton_kernel/linear_att_copy.py | 13 - .../kv_cache_mem_manager/mem_manager.py | 306 ------------ .../qwen3next_mem_manager.py | 455 ------------------ lightllm/common/linear_state.py | 162 +++---- lightllm/common/req_manager/glm5_next.py | 36 -- lightllm/common/req_manager/hybrid_base.py | 72 --- lightllm/common/req_manager/linear_att.py | 55 --- .../common/state_cache_manager/linear_att.py | 205 -------- lightllm/models/glm5_next/config.py | 1 + lightllm/models/glm5_next/model.py | 3 + lightllm/models/qwen3next/model.py | 2 + .../server/router/model_infer/infer_batch.py | 3 - .../model_infer/mtp_speculative/utils.py | 24 +- .../common/att_resources/model_helpers.py | 5 + .../test_hybrid_layer_mapping.py | 14 +- .../att_resources/test_mtp_state_commit.py | 35 +- .../basemodel/attention/linear/test_gdn.py | 13 +- .../basemodel/test_cuda_graph_autotune.py | 7 +- .../linear_att/test_acceptance_fusion.py | 65 ++- .../linear_att/test_mtp_state_params.py | 2 +- .../linear_att/test_replayssm.py | 180 +++---- .../linear_att/test_replayssm_execution.py | 16 +- .../linear_att/test_ssm_autotune.py | 16 +- .../test_hybrid_req_manager_pd_sharing.py | 132 +++++ unit_tests/common/test_linear_state.py | 203 ++++---- .../common/test_linear_state_contract.py | 62 +-- unit_tests/common/test_paged_kv_transfer.py | 45 -- unit_tests/models/glm5_next/test_cache.py | 23 +- unit_tests/models/glm5_next/test_mtp.py | 52 +- unit_tests/models/glm5_next/test_pd_cache.py | 36 +- 37 files changed, 718 insertions(+), 1767 deletions(-) delete mode 100755 lightllm/common/kv_cache_mem_manager/mem_manager.py delete mode 100644 lightllm/common/kv_cache_mem_manager/qwen3next_mem_manager.py delete mode 100644 lightllm/common/req_manager/glm5_next.py delete mode 100644 lightllm/common/req_manager/hybrid_base.py delete mode 100644 lightllm/common/req_manager/linear_att.py delete mode 100644 lightllm/common/state_cache_manager/linear_att.py create mode 100644 unit_tests/common/test_hybrid_req_manager_pd_sharing.py delete mode 100644 unit_tests/common/test_paged_kv_transfer.py diff --git a/lightllm/common/attention_state/bounded_att/gdn/create_gdn_state.py b/lightllm/common/attention_state/bounded_att/gdn/create_gdn_state.py index 3fe6f2317e..fff372e93e 100644 --- a/lightllm/common/attention_state/bounded_att/gdn/create_gdn_state.py +++ b/lightllm/common/attention_state/bounded_att/gdn/create_gdn_state.py @@ -13,6 +13,10 @@ def create_gdn_state( request_capacity: int, mtp_step: int = 0, device: str | torch.device = "cuda", + *, + ssm_state_mode: str = "native", + ssm_replay_config: dict | None = None, + gate_lower_bound: float = -5.0, ) -> GDNState: """创建 GDN 请求运行态，大小页池由后续容量分配流程初始化。""" return GDNState( @@ -20,4 +24,7 @@ def create_gdn_state( request_capacity=request_capacity, mtp_step=mtp_step, device=device, + ssm_state_mode=ssm_state_mode, + ssm_replay_config=ssm_replay_config, + gate_lower_bound=gate_lower_bound, ) diff --git a/lightllm/common/attention_state/bounded_att/gdn/gdn_state.py b/lightllm/common/attention_state/bounded_att/gdn/gdn_state.py index b8a641a73d..fd31143063 100644 --- a/lightllm/common/attention_state/bounded_att/gdn/gdn_state.py +++ b/lightllm/common/attention_state/bounded_att/gdn/gdn_state.py @@ -7,6 +7,7 @@ import torch from lightllm.utils.log_utils import init_logger +from lightllm.common.linear_state import LinearStateManager from ...utils import deepcopy_result, dtype_name from ...types import AxisMapping, StateTransferLayout, TensorField, TransferRegion @@ -23,7 +24,7 @@ logger = init_logger(__name__) -class GDNState(BoundedAttentionState): +class GDNState(LinearStateManager, BoundedAttentionState): """Gated DeltaNet 的 Conv 历史、SSM 和 MTP 游标，支持不同的 K/V head 数及维度。""" big_pages_pool: GDNPagesPool | None @@ -39,35 +40,26 @@ def __init__( request_capacity: int, mtp_step: int = 0, device: str | torch.device = "cuda", + *, + ssm_state_mode: str = "native", + ssm_replay_config: dict | None = None, + gate_lower_bound: float = -5.0, ) -> None: """按请求容量创建 Conv/SSM 运行态和 MTP 游标；大小页池延后初始化。""" assert isinstance(config, ConvSSMStateConfig), "config must be ConvSSMStateConfig" self.config = config self.state_name: str | None = None # 由 AttentionStateManager.register_state() 绑定时设置。 - self.request_capacity = request_capacity # 包含预留的 HOLD 请求槽位。 - self.mtp_step = mtp_step - self._logical_to_storage_layer = { - logical_layer_id: storage_layer_index - for storage_layer_index, logical_layer_id in enumerate(config.logical_layer_ids) - } - storage_layer_count = len(config.logical_layer_ids) - self.conv_state = torch.zeros( - (storage_layer_count, request_capacity, config.conv_channels, config.conv_width + mtp_step), - dtype=config.conv_dtype, - device=device, - ) - self.ssm_state = torch.zeros( - ( - storage_layer_count, - request_capacity * (mtp_step + 1), - config.num_v_heads, - config.key_dim, - config.value_dim, - ), - dtype=config.ssm_dtype, - device=device, + LinearStateManager.__init__( + self, + config, + request_capacity, + mtp_step, + device, + ssm_state_mode=ssm_state_mode, + ssm_replay_config=ssm_replay_config, + kda=False, + gate_lower_bound=gate_lower_bound, ) - self.mtp_state_index = torch.zeros(request_capacity, dtype=torch.int32, device=device) if mtp_step else None self.big_pages_pool = None self.small_pages_pool = None @@ -123,11 +115,7 @@ def reset_bounded_att_state(self, req_indices: List[int]) -> None: """清零指定请求的 Conv 历史、全部 SSM 候选及 MTP 游标，保留页池快照。""" assert isinstance(req_indices, list), "request indices must be a list" for req_idx in req_indices: - self.conv_state[:, req_idx].zero_() - start = req_idx * (self.mtp_step + 1) - self.ssm_state[:, start : start + self.mtp_step + 1].zero_() - if self.mtp_state_index is not None: - self.mtp_state_index[req_idx].fill_(0) + self.reset(req_idx) def save_bounded_att_runtime_to_pages( self, @@ -142,12 +130,11 @@ def save_bounded_att_runtime_to_pages( ), "request and slot indices must be lists" assert len(req_indices) == len(state_token_lens) == len(page_slot_ids), "request, length and page counts differ" pool = self._pool(page_kind) - # TODO：大小页仅在 prefill 阶段保存，mtp_state_index 始终为 0，可直接复制基准状态。 - # 若以后支持 decode/MTP 后保存，需要按已接受游标选择 Conv 窗口和 SSM 槽位。 conv_copies, ssm_copies = [], [] for req_idx, slot_id in zip(req_indices, page_slot_ids): - conv_copies.append((self.conv_state[:, req_idx, :, : self.config.conv_width], pool.conv_state[slot_id])) - ssm_copies.append((self.ssm_state[:, req_idx * (self.mtp_step + 1)], pool.ssm_state[slot_id])) + accepted = self.get_accepted_state(req_idx) + conv_copies.append((accepted.conv, pool.conv_state[slot_id])) + ssm_copies.append((accepted.ssm, pool.ssm_state[slot_id])) copy_bounded_att_tensors(copies=conv_copies) copy_bounded_att_tensors(copies=ssm_copies) @@ -166,8 +153,9 @@ def restore_bounded_att_runtime_from_pages( pool = self._pool(page_kind) conv_copies, ssm_copies = [], [] for req_idx, slot_id in zip(req_indices, page_slot_ids): - conv_copies.append((pool.conv_state[slot_id], self.conv_state[:, req_idx, :, : self.config.conv_width])) - ssm_copies.append((pool.ssm_state[slot_id], self.ssm_state[:, req_idx * (self.mtp_step + 1)])) + target = self.prepare_restore(req_idx) + conv_copies.append((pool.conv_state[slot_id], target.conv)) + ssm_copies.append((pool.ssm_state[slot_id], target.ssm)) copy_bounded_att_tensors(copies=conv_copies) copy_bounded_att_tensors(copies=ssm_copies) if self.mtp_state_index is not None: @@ -183,6 +171,8 @@ def commit_bounded_att_state( assert req_indices.is_cuda and accepted_steps.is_cuda, "commit inputs must be CUDA tensors" if self.mtp_state_index is not None: self.mtp_state_index[req_indices] = accepted_steps - 1 + if self.ssm_backend is not None: + self.ssm_backend.accept_updates(req_indices, self.mtp_state_index) # ======================================================================== # 三、CPU Cache 数据交换 @@ -301,28 +291,14 @@ def copy_runtime_state_from_dp( assert ( type(remote_dp_bounded_state) is type(self) and remote_dp_bounded_state.config == self.config ), "DP copy requires the same algorithm and equivalent TP state configurations" - # TODO：当前 DP fetch 只复制 prefill 基准状态，源 MTP 游标为 0； - # 若以后支持 decode/MTP 后迁移，需按已接受游标选择源 Conv 窗口及 SSM 槽位。 - # 跨设备源存储必须已通过 peer/IPC 映射，批量复制不改变双方页池。 conv_copies, ssm_copies = [], [] for src_req, dst_req in zip(src_req_indices, dst_req_indices): - conv_copies.append( - ( - remote_dp_bounded_state.conv_state[:, src_req, :, : self.config.conv_width], - self.conv_state[:, dst_req, :, : self.config.conv_width], - ) - ) - ssm_copies.append( - ( - remote_dp_bounded_state.ssm_state[:, src_req * (remote_dp_bounded_state.mtp_step + 1)], - self.ssm_state[:, dst_req * (self.mtp_step + 1)], - ) - ) + source = remote_dp_bounded_state.get_accepted_state(src_req) + target = self.prepare_restore(dst_req) + conv_copies.append((source.conv, target.conv)) + ssm_copies.append((source.ssm, target.ssm)) copy_bounded_att_tensors(copies=conv_copies) copy_bounded_att_tensors(copies=ssm_copies) - if self.mtp_state_index is not None: - for dst_req in dst_req_indices: - self.mtp_state_index[dst_req].fill_(0) # ======================================================================== # 五、PD 分离传输 @@ -342,10 +318,10 @@ def write_to_pd_page( state_token_len: int, layout: StateTransferLayout, ) -> None: - """写出 state_token_len 处的 prefill 基准状态；Conv/SSM 不依赖长度，源 MTP 游标须为 0。""" + """写出 state_token_len 处的已接受状态，不修改源 Replay 历史。""" assert layout.state_name == self.state_name, "transfer layout does not describe this resource" assert 0 <= layout.byte_start < layout.byte_end, "invalid State byte interval" - fields = self._runtime_fields(req_idx) + fields = self._fields(*self.get_accepted_state(req_idx)) for copies in bounded_att_buffer_copies( page_fields=[fields], page_buffers=pd_page_buffers, @@ -366,7 +342,7 @@ def read_from_pd_page( """恢复 state_token_len 处的基准状态并规范化控制信息；Conv/SSM 不依赖长度定位。""" assert layout.state_name == self.state_name, "transfer layout does not describe this resource" assert 0 <= layout.byte_start < layout.byte_end, "invalid State byte interval" - fields = self._runtime_fields(req_idx) + fields = self._fields(*self.prepare_restore(req_idx)) for copies in bounded_att_buffer_copies( page_fields=[fields], page_buffers=pd_page_buffers, @@ -506,14 +482,6 @@ def _pool(self, page_kind: PageKind): assert pool is not None, f"{page_kind} pages are disabled for {self.state_name}" return pool - def _runtime_fields(self, req: int): - """返回请求基准位置的 Conv/SSM 视图，供恢复和传输使用。""" - # TODO：当前 DP/PD 不会发生在 MTP 提交后，源状态游标为 0。 - # 若以后支持该场景，需要按已接受游标选取 Conv 窗口及 SSM 槽位。 - return self._fields( - self.conv_state[:, req, :, : self.config.conv_width], self.ssm_state[:, req * (self.mtp_step + 1)] - ) - @staticmethod def _fields(conv: torch.Tensor, ssm: torch.Tensor): return { diff --git a/lightllm/common/attention_state/bounded_att/kda/create_kda_state.py b/lightllm/common/attention_state/bounded_att/kda/create_kda_state.py index 49c9ad0ce5..180c5dcc7b 100644 --- a/lightllm/common/attention_state/bounded_att/kda/create_kda_state.py +++ b/lightllm/common/attention_state/bounded_att/kda/create_kda_state.py @@ -13,6 +13,10 @@ def create_kda_state( request_capacity: int, mtp_step: int = 0, device: str | torch.device = "cuda", + *, + ssm_state_mode: str = "native", + ssm_replay_config: dict | None = None, + gate_lower_bound: float = -5.0, ) -> KDAState: """创建 KDA 请求运行态，大小页池由后续容量分配流程初始化。""" return KDAState( @@ -20,4 +24,7 @@ def create_kda_state( request_capacity=request_capacity, mtp_step=mtp_step, device=device, + ssm_state_mode=ssm_state_mode, + ssm_replay_config=ssm_replay_config, + gate_lower_bound=gate_lower_bound, ) diff --git a/lightllm/common/attention_state/bounded_att/kda/kda_state.py b/lightllm/common/attention_state/bounded_att/kda/kda_state.py index e5480a3ab0..9f8051836f 100644 --- a/lightllm/common/attention_state/bounded_att/kda/kda_state.py +++ b/lightllm/common/attention_state/bounded_att/kda/kda_state.py @@ -7,6 +7,7 @@ import torch from lightllm.utils.log_utils import init_logger +from lightllm.common.linear_state import LinearStateManager from ...utils import deepcopy_result, dtype_name from ...types import AxisMapping, StateTransferLayout, TensorField, TransferRegion @@ -23,7 +24,7 @@ logger = init_logger(__name__) -class KDAState(BoundedAttentionState): +class KDAState(LinearStateManager, BoundedAttentionState): """KDA 的 Conv 历史、方形 SSM 和 MTP 游标，要求 Q/K/V head 数及维度相同。""" big_pages_pool: KDAPagesPool | None @@ -39,35 +40,26 @@ def __init__( request_capacity: int, mtp_step: int = 0, device: str | torch.device = "cuda", + *, + ssm_state_mode: str = "native", + ssm_replay_config: dict | None = None, + gate_lower_bound: float = -5.0, ) -> None: """按请求容量创建 Conv/SSM 运行态和 MTP 游标；大小页池延后初始化。""" assert isinstance(config, ConvSSMStateConfig), "config must be ConvSSMStateConfig" self.config = config self.state_name: str | None = None # 由 AttentionStateManager.register_state() 绑定时设置。 - self.request_capacity = request_capacity # 包含预留的 HOLD 请求槽位。 - self.mtp_step = mtp_step - self._logical_to_storage_layer = { - logical_layer_id: storage_layer_index - for storage_layer_index, logical_layer_id in enumerate(config.logical_layer_ids) - } - storage_layer_count = len(config.logical_layer_ids) - self.conv_state = torch.zeros( - (storage_layer_count, request_capacity, config.conv_channels, config.conv_width + mtp_step), - dtype=config.conv_dtype, - device=device, - ) - self.ssm_state = torch.zeros( - ( - storage_layer_count, - request_capacity * (mtp_step + 1), - config.num_v_heads, - config.key_dim, - config.value_dim, - ), - dtype=config.ssm_dtype, - device=device, + LinearStateManager.__init__( + self, + config, + request_capacity, + mtp_step, + device, + ssm_state_mode=ssm_state_mode, + ssm_replay_config=ssm_replay_config, + kda=True, + gate_lower_bound=gate_lower_bound, ) - self.mtp_state_index = torch.zeros(request_capacity, dtype=torch.int32, device=device) if mtp_step else None self.big_pages_pool = None self.small_pages_pool = None @@ -123,11 +115,7 @@ def reset_bounded_att_state(self, req_indices: List[int]) -> None: """清零指定请求的 Conv 历史、全部 SSM 候选及 MTP 游标，保留页池快照。""" assert isinstance(req_indices, list), "request indices must be a list" for req_idx in req_indices: - self.conv_state[:, req_idx].zero_() - start = req_idx * (self.mtp_step + 1) - self.ssm_state[:, start : start + self.mtp_step + 1].zero_() - if self.mtp_state_index is not None: - self.mtp_state_index[req_idx].fill_(0) + self.reset(req_idx) def save_bounded_att_runtime_to_pages( self, @@ -142,12 +130,11 @@ def save_bounded_att_runtime_to_pages( ), "request and slot indices must be lists" assert len(req_indices) == len(state_token_lens) == len(page_slot_ids), "request, length and page counts differ" pool = self._pool(page_kind) - # TODO：大小页仅在 prefill 阶段保存，mtp_state_index 始终为 0，可直接复制基准状态。 - # 若以后支持 decode/MTP 后保存，需要按已接受游标选择 Conv 窗口和 SSM 槽位。 conv_copies, ssm_copies = [], [] for req_idx, slot_id in zip(req_indices, page_slot_ids): - conv_copies.append((self.conv_state[:, req_idx, :, : self.config.conv_width], pool.conv_state[slot_id])) - ssm_copies.append((self.ssm_state[:, req_idx * (self.mtp_step + 1)], pool.ssm_state[slot_id])) + accepted = self.get_accepted_state(req_idx) + conv_copies.append((accepted.conv, pool.conv_state[slot_id])) + ssm_copies.append((accepted.ssm, pool.ssm_state[slot_id])) copy_bounded_att_tensors(copies=conv_copies) copy_bounded_att_tensors(copies=ssm_copies) @@ -166,8 +153,9 @@ def restore_bounded_att_runtime_from_pages( pool = self._pool(page_kind) conv_copies, ssm_copies = [], [] for req_idx, slot_id in zip(req_indices, page_slot_ids): - conv_copies.append((pool.conv_state[slot_id], self.conv_state[:, req_idx, :, : self.config.conv_width])) - ssm_copies.append((pool.ssm_state[slot_id], self.ssm_state[:, req_idx * (self.mtp_step + 1)])) + target = self.prepare_restore(req_idx) + conv_copies.append((pool.conv_state[slot_id], target.conv)) + ssm_copies.append((pool.ssm_state[slot_id], target.ssm)) copy_bounded_att_tensors(copies=conv_copies) copy_bounded_att_tensors(copies=ssm_copies) if self.mtp_state_index is not None: @@ -183,6 +171,8 @@ def commit_bounded_att_state( assert req_indices.is_cuda and accepted_steps.is_cuda, "commit inputs must be CUDA tensors" if self.mtp_state_index is not None: self.mtp_state_index[req_indices] = accepted_steps - 1 + if self.ssm_backend is not None: + self.ssm_backend.accept_updates(req_indices, self.mtp_state_index) # ======================================================================== # 三、CPU Cache 数据交换 @@ -301,28 +291,14 @@ def copy_runtime_state_from_dp( assert ( type(remote_dp_bounded_state) is type(self) and remote_dp_bounded_state.config == self.config ), "DP copy requires the same algorithm and equivalent TP state configurations" - # TODO：当前 DP fetch 只复制 prefill 基准状态，源 MTP 游标为 0； - # 若以后支持 decode/MTP 后迁移，需按已接受游标选择源 Conv 窗口及 SSM 槽位。 - # 跨设备源存储必须已通过 peer/IPC 映射，批量复制不改变双方页池。 conv_copies, ssm_copies = [], [] for src_req, dst_req in zip(src_req_indices, dst_req_indices): - conv_copies.append( - ( - remote_dp_bounded_state.conv_state[:, src_req, :, : self.config.conv_width], - self.conv_state[:, dst_req, :, : self.config.conv_width], - ) - ) - ssm_copies.append( - ( - remote_dp_bounded_state.ssm_state[:, src_req * (remote_dp_bounded_state.mtp_step + 1)], - self.ssm_state[:, dst_req * (self.mtp_step + 1)], - ) - ) + source = remote_dp_bounded_state.get_accepted_state(src_req) + target = self.prepare_restore(dst_req) + conv_copies.append((source.conv, target.conv)) + ssm_copies.append((source.ssm, target.ssm)) copy_bounded_att_tensors(copies=conv_copies) copy_bounded_att_tensors(copies=ssm_copies) - if self.mtp_state_index is not None: - for dst_req in dst_req_indices: - self.mtp_state_index[dst_req].fill_(0) # ======================================================================== # 五、PD 分离传输 @@ -342,10 +318,10 @@ def write_to_pd_page( state_token_len: int, layout: StateTransferLayout, ) -> None: - """写出 state_token_len 处的 prefill 基准状态；Conv/SSM 不依赖长度，源 MTP 游标须为 0。""" + """写出 state_token_len 处的已接受状态，不修改源 Replay 历史。""" assert layout.state_name == self.state_name, "transfer layout does not describe this resource" assert 0 <= layout.byte_start < layout.byte_end, "invalid State byte interval" - fields = self._runtime_fields(req_idx) + fields = self._fields(*self.get_accepted_state(req_idx)) for copies in bounded_att_buffer_copies( page_fields=[fields], page_buffers=pd_page_buffers, @@ -366,7 +342,7 @@ def read_from_pd_page( """恢复 state_token_len 处的基准状态并规范化控制信息；Conv/SSM 不依赖长度定位。""" assert layout.state_name == self.state_name, "transfer layout does not describe this resource" assert 0 <= layout.byte_start < layout.byte_end, "invalid State byte interval" - fields = self._runtime_fields(req_idx) + fields = self._fields(*self.prepare_restore(req_idx)) for copies in bounded_att_buffer_copies( page_fields=[fields], page_buffers=pd_page_buffers, @@ -506,14 +482,6 @@ def _pool(self, page_kind: PageKind): assert pool is not None, f"{page_kind} pages are disabled for {self.state_name}" return pool - def _runtime_fields(self, req: int): - """返回请求基准位置的 Conv/SSM 视图，供恢复和传输使用。""" - # TODO：当前 DP/PD 不会发生在 MTP 提交后，源状态游标为 0。 - # 若以后支持该场景，需要按已接受游标选取 Conv 窗口及 SSM 槽位。 - return self._fields( - self.conv_state[:, req, :, : self.config.conv_width], self.ssm_state[:, req * (self.mtp_step + 1)] - ) - @staticmethod def _fields(conv: torch.Tensor, ssm: torch.Tensor): return { diff --git a/lightllm/common/basemodel/attention/linear/gdn.py b/lightllm/common/basemodel/attention/linear/gdn.py index 7c6d2939a2..32403a379f 100644 --- a/lightllm/common/basemodel/attention/linear/gdn.py +++ b/lightllm/common/basemodel/attention/linear/gdn.py @@ -1,4 +1,5 @@ from lightllm.common.attention_state.bounded_att import GDNState +from lightllm.common.linear_state import get_linear_state import dataclasses import weakref from abc import ABC, abstractmethod @@ -120,7 +121,7 @@ def init_state(self): # 每次 _prefill 都会在 runtime infer_state 上调用 init_state。 # prefill cuda graph 回调必须走 new_infer_state.prefill_att_state1， # 才能读到这里按当前 batch（含 token padding 后的 dummy request）更新的索引。 - linear_state_manager = self.infer_state.req_manager.linear_state_manager + linear_state_manager = get_linear_state(self.infer_state, GDNState) self.b_conv_buffer_idx, self.b_ssm_buffer_idx = linear_state_manager.prepare_prefill(self.infer_state.b_req_idx) return @@ -141,7 +142,7 @@ def prefill_att( layer_num = linear_att_dict["layer_num"] backend: GDNLinearAttBackend = self.backend - conv_states, ssm_states = self.infer_state.req_manager.get_mamba_cache(layer_num) + conv_states, ssm_states = get_linear_state(self.infer_state, GDNState).get_layer_view(layer_num) # 在开启了mtp的时候，conv 状态的最后一维可能存在冗余的部分，需要进行切片对齐。 # prefill 模式下，使用不到这几个维度，所以需要扣除掉， if backend.mtp_step > 0: @@ -210,7 +211,7 @@ class GDNLinearAttDecodeAttState(BaseDecodeAttState): b_num_accepted_tokens: torch.Tensor = None def get_linear_state_commit(self) -> Optional["LinearStateCommit"]: - return self.infer_state.req_manager.linear_state_manager.create_commit(self.run_config) + return get_linear_state(self.infer_state, GDNState).create_commit(self.run_config) def init_state(self): draft_step = self.backend.model.mtp_manager.get_decode_draft_step(self.backend.model.is_mtp_draft_model) @@ -220,7 +221,7 @@ def init_state(self): self._init_dynamic_mtp_decode_state() else: self._init_fixed_mtp_decode_state(draft_step) - linear_state_manager = self.infer_state.req_manager.linear_state_manager + linear_state_manager = get_linear_state(self.infer_state, GDNState) if draft_step == 0: linear_state_manager.prepare_decode(self.b_conv_buffer_idx, single_token=True) else: @@ -241,7 +242,7 @@ def _init_dynamic_mtp_decode_state(self): hold_req_id=self.infer_state.req_manager.HOLD_REQUEST_ID, ) - def _init_fixed_mtp_decode_state(self, draft_step: int, *, req_to_mtp_state_index: torch.Tensor) -> None: + def _init_fixed_mtp_decode_state(self, draft_step: int) -> None: mtp_size = draft_step + 1 batch_size = self.infer_state.batch_size assert batch_size % mtp_size == 0, ( @@ -277,7 +278,7 @@ def decode_att( backend: GDNLinearAttBackend = self.backend mixed_qkv, z, b, a = backend._split_qkvzba(mixed_qkvzba) - conv_states, ssm_states = self.infer_state.req_manager.get_mamba_cache(layer_num) + conv_states, ssm_states = get_linear_state(self.infer_state, GDNState).get_layer_view(layer_num) draft_step = self.backend.model.mtp_manager.get_decode_draft_step(self.backend.model.is_mtp_draft_model) if draft_step > 0: @@ -336,9 +337,9 @@ def _gdn_decode_kernel( backend.tp_num_v_heads, backend.head_v_dim, ) - ssm_backend = self.infer_state.req_manager.linear_state_manager.ssm_backend + ssm_backend = get_linear_state(self.infer_state, GDNState).ssm_backend if ssm_backend is not None: - linear_layer_idx = self.infer_state.req_manager.linear_config.get_linear_layer_index( + linear_layer_idx = get_linear_state(self.infer_state, GDNState).get_storage_layer_index( layer_weight.layer_num_ ) core_attn_out, self.run_config = ssm_backend.forward( @@ -399,9 +400,9 @@ def _gdn_mtp_kernel( ) query, key, value = backend._rearrange_mixed_qkv(mixed_qkv, decode=False) - ssm_backend = self.infer_state.req_manager.linear_state_manager.ssm_backend + ssm_backend = get_linear_state(self.infer_state, GDNState).ssm_backend if ssm_backend is not None: - linear_layer_idx = self.infer_state.req_manager.linear_config.get_linear_layer_index( + linear_layer_idx = get_linear_state(self.infer_state, GDNState).get_storage_layer_index( layer_weight.layer_num_ ) core_attn_out, self.run_config = ssm_backend.forward( diff --git a/lightllm/common/basemodel/attention/linear/kda.py b/lightllm/common/basemodel/attention/linear/kda.py index 3e7a7634b6..2fa1cbbba6 100644 --- a/lightllm/common/basemodel/attention/linear/kda.py +++ b/lightllm/common/basemodel/attention/linear/kda.py @@ -11,6 +11,7 @@ import torch from lightllm.common.attention_state.bounded_att import KDAState +from lightllm.common.linear_state import get_linear_state from lightllm.common.basemodel.attention.base_att import ( AttControl, BaseAttBackend, @@ -65,7 +66,7 @@ class KDALinearAttPrefillAttState(BasePrefillAttState): b_ssm_buffer_idx: torch.Tensor = None def init_state(self): - linear_state_manager = self.infer_state.req_manager.linear_state_manager + linear_state_manager = get_linear_state(self.infer_state, KDAState) self.b_conv_buffer_idx, self.b_ssm_buffer_idx = linear_state_manager.prepare_prefill(self.infer_state.b_req_idx) def prefill_att( @@ -140,7 +141,7 @@ class KDALinearAttDecodeAttState(BaseDecodeAttState): b_num_accepted_tokens: torch.Tensor = None def get_linear_state_commit(self) -> Optional["LinearStateCommit"]: - return self.infer_state.req_manager.linear_state_manager.create_commit(self.run_config) + return get_linear_state(self.infer_state, KDAState).create_commit(self.run_config) def init_state(self): mtp_step = self.backend.mtp_step @@ -150,7 +151,7 @@ def init_state(self): self._init_dynamic_mtp_decode_state() else: self._init_fixed_mtp_decode_state(mtp_step) - linear_state_manager = self.infer_state.req_manager.linear_state_manager + linear_state_manager = get_linear_state(self.infer_state, KDAState) if mtp_step == 0: linear_state_manager.prepare_decode(self.b_conv_buffer_idx, single_token=True) else: @@ -171,7 +172,7 @@ def _init_dynamic_mtp_decode_state(self): hold_req_id=self.infer_state.req_manager.HOLD_REQUEST_ID, ) - def _init_fixed_mtp_decode_state(self, mtp_step: int, *, req_to_mtp_state_index: torch.Tensor) -> None: + def _init_fixed_mtp_decode_state(self, mtp_step: int) -> None: mtp_size = mtp_step + 1 batch_size = self.infer_state.batch_size assert batch_size % mtp_size == 0, ( @@ -230,9 +231,9 @@ def decode_att( v = v.view(*shape, backend.tp_num_heads, backend.head_dim) raw_gate = raw_gate.view(*shape, backend.tp_hidden_size) raw_beta = raw_beta.view(*shape, backend.tp_num_heads) - ssm_backend = self.infer_state.req_manager.linear_state_manager.ssm_backend + ssm_backend = get_linear_state(self.infer_state, KDAState).ssm_backend if ssm_backend is not None: - linear_layer_idx = self.infer_state.req_manager.linear_config.get_linear_layer_index(layer_num) + linear_layer_idx = get_linear_state(self.infer_state, KDAState).get_storage_layer_index(layer_num) output, self.run_config = ssm_backend.forward( linear_layer_idx, q, diff --git a/lightllm/common/basemodel/basemodel.py b/lightllm/common/basemodel/basemodel.py index 2c42547659..af2ae932be 100755 --- a/lightllm/common/basemodel/basemodel.py +++ b/lightllm/common/basemodel/basemodel.py @@ -5,8 +5,6 @@ import copy import math import json -from typing import final, List - import torch import torch.nn.functional as F import triton diff --git a/lightllm/common/basemodel/triton_kernel/linear_att_copy.py b/lightllm/common/basemodel/triton_kernel/linear_att_copy.py index 45568703b6..eb17507fb0 100644 --- a/lightllm/common/basemodel/triton_kernel/linear_att_copy.py +++ b/lightllm/common/basemodel/triton_kernel/linear_att_copy.py @@ -25,9 +25,6 @@ def _copy_linear_att_state_to_kv_buffer( cpu_kv_ssm_stride_l, cpu_kv_ssm_stride_d, mtp_step, - accepted_state_indices, - conv_element_bytes: tl.constexpr, - HAS_ACCEPTED_STATE_INDICES: 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, @@ -56,11 +53,6 @@ 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_ACCEPTED_STATE_INDICES: - accepted_offset = tl.load(accepted_state_indices + cur_req_idx) - gpu_conv_base += accepted_offset * conv_element_bytes - if mtp_step > 0: - cur_state_req_idx += accepted_offset 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)): @@ -94,12 +86,10 @@ 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, - accepted_state_indices: 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]" @@ -154,9 +144,6 @@ 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, - accepted_state_indices=accepted_state_indices, - conv_element_bytes=conv_element_bytes, - HAS_ACCEPTED_STATE_INDICES=accepted_state_indices 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/mem_manager.py b/lightllm/common/kv_cache_mem_manager/mem_manager.py deleted file mode 100755 index f606ceb30b..0000000000 --- a/lightllm/common/kv_cache_mem_manager/mem_manager.py +++ /dev/null @@ -1,306 +0,0 @@ -import math -import os -import re -from typing import List, Tuple, Any, Union -from multiprocessing.reduction import ForkingPickler - -import torch -import torch.distributed as dist -import torch.multiprocessing as mp -from filelock import FileLock - -from lightllm.utils.log_utils import init_logger -from lightllm.server.router.dynamic_prompt.shared_arr import SharedInt -from .allocator import KvCacheAllocator -from lightllm.utils.profile_max_tokens import get_available_gpu_memory, get_total_gpu_memory -from lightllm.utils.dist_utils import ( - get_current_device_id, - get_current_rank_in_node, - get_node_world_size, -) -from lightllm.utils.envs_utils import get_env_start_args -from lightllm.utils.shm_utils import get_service_shm_name -from lightllm.utils.config_utils import get_num_key_value_heads -from lightllm.common.kv_trans_kernel.nixl_kv_trans import page_io -from lightllm.utils.device_utils import kv_trans_use_p2p -from lightllm.utils.shm_utils import create_or_link_shm -from .operator import BaseMemManagerOperator, NormalMemOperator - -logger = init_logger(__name__) - - -class MemoryManager: - - operator_class = NormalMemOperator - - def __init__(self, size, dtype, head_num, head_dim, layer_num, always_copy=False, mem_fraction=0.9): - self.size = size - self.head_num = head_num - self.head_dim = head_dim - self.layer_num = layer_num - self.always_copy = always_copy - self.dtype = dtype - # profile the max total token num if the size is None - self.profile_size(mem_fraction) - - # A physical KV page must never straddle the allocator boundary. The - # unused remainder is intentionally dropped so every managed page has - # a stable ``mem_index // page_size`` id. - self.page_size = get_env_start_args().page_size - self.size = self.size // self.page_size * self.page_size - - self.allocator = KvCacheAllocator(self.size) - - self._init_buffers( - self.size, - dtype, - head_num, - head_dim, - layer_num, - ) - self.HOLD_TOKEN_MEMINDEXES = tuple(range(self.size, self.size + self.page_size)) - - # 构建对外的操作类接口 - self.operator: BaseMemManagerOperator = self.operator_class(self) - - def get_att_input_params(self, layer_index: int) -> Tuple[Any, Any]: - k = self.kv_buffer[layer_index][:, : self.head_num, :] - v = self.kv_buffer[layer_index][:, self.head_num :, :] - return k, v - - def get_cell_size(self): - return 2 * self.head_num * self.head_dim * self.layer_num * torch._utils._element_size(self.dtype) - - def get_paged_kv_move_buffer_shape(self, page_num, page_size): - num_kv_head = get_num_key_value_heads(get_env_start_args().model_dir) - return ( - page_num, - page_size, - self.layer_num, - 2 * num_kv_head, - self.head_dim, - ) - - def get_pd_kv_move_buffer_size(self): - args = get_env_start_args() - if args.run_mode not in ["prefill", "decode"]: - return 0 - shape = self.get_paged_kv_move_buffer_shape( - page_num=args.pd_kv_page_num, - page_size=args.pd_kv_page_size, - ) - return math.prod(shape) * torch._utils._element_size(self.dtype) - - def profile_size(self, mem_fraction): - if self.size is not None: - return - - torch.cuda.empty_cache() - world_size = dist.get_world_size() - available_memory = get_available_gpu_memory(world_size) - get_total_gpu_memory() * (1 - mem_fraction) - cell_size = self.get_cell_size() - pd_kv_move_buffer_size = self.get_pd_kv_move_buffer_size() - available_memory_bytes = available_memory * 1024 ** 3 - pd_kv_move_buffer_size - self.size = int(available_memory_bytes / cell_size) - if world_size > 1: - tensor = torch.tensor(self.size, dtype=torch.int64, device=f"cuda:{get_current_device_id()}") - dist.all_reduce(tensor, op=dist.ReduceOp.MIN) - self.size = tensor.item() - logger.info( - f"{str(available_memory)} GB space is available after load the model weight\n" - f"{str(pd_kv_move_buffer_size / 1024 ** 2)} MB is reserved for PD KV transfer buffer\n" - f"{str(cell_size / 1024 ** 2)} MB is the size of one token kv cache\n" - f"{self.size} is the profiled max_total_token_num with the mem_fraction {mem_fraction}\n" - ) - return - - def _init_buffers(self, size, dtype, head_num, head_dim, layer_num): - assert size % self.page_size == 0, f"KV cache size {size} must be a multiple of page_size {self.page_size}" - # 在初始化 kv_buffer 时，每层额外保留一整页；这部分空间不会被 allocator 对外分配。 - # 这一页用于多 DP、overlap microbatch 等模式下的 padding 请求， - # 页内索引存储在 HOLD_TOKEN_MEMINDEXES 中，与 req_manager 的 HOLD_REQUEST_ID 作用类似。 - self.kv_buffer = torch.empty( - (layer_num, size + self.page_size, 2 * head_num, head_dim), dtype=dtype, device="cuda" - ) - - def alloc_paged_kv_move_buffer(self, page_num, page_size) -> torch.Tensor: - self.kv_move_buffer = torch.empty( - self.get_paged_kv_move_buffer_shape(page_num, page_size), dtype=self.dtype, device="cuda" - ) - self._buffer_mem_indexes_tensors = [ - torch.empty((page_size,), dtype=torch.int64, device="cpu", pin_memory=True) for _ in range(page_num) - ] - return self.kv_move_buffer - - def write_mem_to_page_kv_move_buffer( - self, - mem_indexes: List[int], - page_index: int, - dp_index: int, - mem_managers: List["MemoryManager"], - dp_world_size: int, - page_kind: str = "kv", - req_idx: int = None, - ): - assert page_kind == "kv", f"{type(self).__name__} does not support page_kind={page_kind}" - cur_page = self.kv_move_buffer[page_index] - pin_mem_indexes = self._buffer_mem_indexes_tensors[page_index][0 : len(mem_indexes)] - pin_mem_indexes.numpy()[:] = mem_indexes - mem_indexes_gpu = pin_mem_indexes.cuda(non_blocking=True) - repeat_count = dp_world_size * self.kv_buffer.shape[2] // self.kv_move_buffer.shape[3] - dp_mems = mem_managers[(dp_index * dp_world_size) : ((dp_index + 1) * dp_world_size)] - for tp_index in range(dp_world_size): - if tp_index % repeat_count == 0: - page_io( - mem_indexes=mem_indexes_gpu, - page_tensor=cur_page, - kv_buffer=dp_mems[tp_index].kv_buffer, - tp_index=tp_index, - tp_world_size=dp_world_size, - mode="write", - ) - # keep for debug - # logger.info(f"src token tensor {self.kv_buffer[:, mem_indexes[0], 0, 0]}") - # logger.info(f"src page token tensor {cur_page[0, :, 0, 0]}") - return - - def read_page_kv_move_buffer_to_mem( - self, - mem_indexes: List[int], - page_index: int, - dp_index: int, - mem_managers: List["MemoryManager"], - dp_world_size: int, - page_kind: str = "kv", - req_idx: int = None, - ): - assert page_kind == "kv", f"{type(self).__name__} does not support page_kind={page_kind}" - cur_page = self.kv_move_buffer[page_index] - pin_mem_indexes = self._buffer_mem_indexes_tensors[page_index][0 : len(mem_indexes)] - pin_mem_indexes.numpy()[:] = mem_indexes - mem_indexes_gpu = pin_mem_indexes.cuda(non_blocking=True) - dp_mems = mem_managers[(dp_index * dp_world_size) : ((dp_index + 1) * dp_world_size)] - for tp_index in range(dp_world_size): - page_io( - mem_indexes=mem_indexes_gpu, - page_tensor=cur_page, - kv_buffer=dp_mems[tp_index].kv_buffer, - tp_index=tp_index, - tp_world_size=dp_world_size, - mode="read", - ) - # keep for debug - # logger.info(f"dst token tensor {self.kv_buffer[:, mem_indexes[0], 0, 0]}") - # logger.info(f"dst page token tensor {cur_page[0, :, 0, 0]}") - - def _free_buffers(self): - self.kv_buffer = None - - def alloc(self, need_size) -> torch.Tensor: - assert ( - need_size % self.page_size == 0 - ), f"KV cache allocation size {need_size} must be a multiple of page_size {self.page_size}" - return self.allocator.alloc(need_size) - - def free(self, free_index: Union[torch.Tensor, List[int]]) -> None: - self.allocator.free(free_index) - - def free_all(self): - self.allocator.free_all() - - def resize_mem(self, new_size): - """ - just for test code - """ - size = new_size // self.page_size * self.page_size - dtype = self.dtype - head_num = self.head_num - head_dim = self.head_dim - layer_num = self.layer_num - - self.size = size - self.allocator.resize(size) - self._free_buffers() - self._init_buffers(size, dtype, head_num, head_dim, layer_num) - self.HOLD_TOKEN_MEMINDEXES = tuple(range(self.size, self.size + self.page_size)) - return - - def get_index_kv_buffer(self, index): - return {"kv_buffer": self.kv_buffer[:, index]} - - def load_index_kv_buffer(self, index, load_tensor_dict): - self.kv_buffer[:, index].copy_(load_tensor_dict["kv_buffer"]) - - def write_to_shm(self, req_manager): - """ - 将 mem manager 写入到 shm中，方便pd分离等特性直接从中读取，不依赖进程间队列。 - """ - if kv_trans_use_p2p(): - from lightllm.server.router.model_infer.mode_backend.pd.p2p_fix import reduce_tensor - - mp.reductions.reduce_tensor.__code__ = reduce_tensor.__code__ - - from lightllm.common.req_manager import ReqManager - - req_manager: ReqManager = req_manager - - # 这个地方是一个不太优雅的设计，但是暂时这么做，可以让dp shared kv swap模块直接访问 req_manager 中的 req_to_token_indexs - # 避免过多无用的数据复制和传输开销。 - self.req_to_token_indexs: torch.Tensor = req_manager.req_to_token_indexs - - lock = FileLock(f"/tmp/{get_service_shm_name('mem_manager_lock')}") - with lock: - node_world_size = get_node_world_size() - shm_name = f"mem_manager_{get_current_rank_in_node()}" - obj_bytes_array = [ForkingPickler.dumps(self).tobytes() for _ in range(node_world_size * 2)] - obj_size = len(obj_bytes_array[0]) - shm = create_or_link_shm( - name=shm_name, expected_size=obj_size * (node_world_size * 2) + 4 + 4, force_mode="create" - ) - logger.info(f"create shm {shm.name} size {shm.size} for mem manger shared buffer") - shm.buf[0:4] = (node_world_size * 2).to_bytes(4, "little") - shm.buf[4:8] = obj_size.to_bytes(4, "little") - start_index = 8 - for obj_bytes in obj_bytes_array: - shm.buf[start_index : start_index + obj_size] = obj_bytes - start_index += obj_size - - @staticmethod - def loads_from_shm(rank_in_node: int) -> "MemoryManager": - shm_name = f"mem_manager_{rank_in_node}" - lock = FileLock(f"/tmp/{get_service_shm_name('mem_manager_lock')}") - logger.info(f"get memmanager from shm {shm_name}") - with lock: - shm = create_or_link_shm(name=shm_name, expected_size=-1, force_mode="link") - left_num = int.from_bytes(shm.buf[0:4], "little") - obj_size = int.from_bytes(shm.buf[4:8], "little") - assert left_num > 0 - end_index = 8 + left_num * obj_size - start_index = 8 + (left_num - 1) * obj_size - obj_bytes = shm.buf[start_index:end_index].tobytes() - shm.buf[0:4] = (left_num - 1).to_bytes(4, byteorder="little") - shm.close() - return ForkingPickler.loads(obj_bytes) - - -class ReadOnlyStaticsMemoryManager: - """ - 读取一些统计信息 - """ - - def __init__(self) -> None: - args = get_env_start_args() - self.global_world_size = args.tp - self.node_world_size = args.tp // args.nnodes - self.dp_world_size = self.global_world_size // args.dp - # 兼容多机 dp size=1 纯 tp 模式的情况 - self.is_multinode_tp = args.dp == 1 and args.nnodes > 1 - self.shared_tp_infos = [ - SharedInt(f"mem_manger_can_use_token_num_{rank_in_node}") - for rank_in_node in range(0, self.node_world_size, self.dp_world_size) - ] - - def get_unrefed_token_num(self, dp_rank_in_node: int): - if self.is_multinode_tp: - return self.shared_tp_infos[0].get_value() - return self.shared_tp_infos[dp_rank_in_node].get_value() diff --git a/lightllm/common/kv_cache_mem_manager/qwen3next_mem_manager.py b/lightllm/common/kv_cache_mem_manager/qwen3next_mem_manager.py deleted file mode 100644 index 0320632015..0000000000 --- a/lightllm/common/kv_cache_mem_manager/qwen3next_mem_manager.py +++ /dev/null @@ -1,455 +0,0 @@ -from typing import Tuple, Any, List - -import torch -import triton -from lightllm.utils.log_utils import init_logger -from lightllm.common.kv_cache_mem_manager.mem_manager import MemoryManager -from lightllm.utils.envs_utils import get_env_start_args -from lightllm.common.state_cache_manager import LinearAttCacheConfig, LinearAttCacheManager -from .fp8_static_per_head_quant_mem_manager import FP8StaticPerHeadQuantMemManager -from .fp8_static_per_tensor_quant_mem_manager import FP8StaticPerTensorQuantMemManager -from .operator import ( - FP8StaticPerHeadQuantMemOperator, - FP8StaticPerTensorQuantMemOperator, - LinearAttMemOperator, -) - -logger = init_logger(__name__) - - -class Qwen3NextMemManager(MemoryManager): - operator_class = LinearAttMemOperator - - def __init__( - self, - size, - dtype, - num_kv_heads, - head_dim, - full_att_layer_num, - linear_config: LinearAttCacheConfig, - always_copy=False, - mem_fraction=0.9, - ): - self.linear_config = linear_config - - super().__init__(size, dtype, num_kv_heads, head_dim, full_att_layer_num, always_copy, mem_fraction) - - def get_att_input_params(self, layer_index: int) -> Tuple[Any, Any]: - layer_index = self.linear_config.get_full_att_kv_layer_index(layer_index) - return super().get_att_input_params(layer_index) - - def _init_buffers(self, size, dtype, head_num, head_dim, layer_num): - super()._init_buffers(size, dtype, head_num, head_dim, layer_num) - # TODO 初始化线性 att 对应的部分 buffer. - self._init_linear_att_buffers() - return - - def _init_linear_att_buffers(self): - big_page_token_num = ( - get_env_start_args().linear_att_page_block_num * get_env_start_args().linear_att_hash_page_size - ) - # 申请大页可能需要对应的资源, 多申请了两个linear att的状态，理论上这个状态 - # 永远不会被 alloc 申请到，只会在 cpu cache中，用于过渡和存储碎页情况下的 - # cpu cache 的页面拷贝。 - self.big_page_buffers = LinearAttCacheManager( - size=triton.cdiv(self.size, big_page_token_num) + 2, - linear_config=self.linear_config, - keep_num=2, - ) - - self.CPU_CACHE_BIG_PAGE_LOAD_TEMP_BUFFER_ID = self.big_page_buffers.size - 2 - self.CPU_CACHE_BIG_PAGE_OFFLOAD_TEMP_BUFFER_ID = self.big_page_buffers.size - 1 - return - - def _free_buffers(self): - super()._free_buffers() - self._free_linear_att_buffers() - return - - def _free_linear_att_buffers(self): - self.big_page_buffers = None - return - - def write_to_shm(self, req_manager): - # LinearStateManager contains only IPC-safe state/operator configuration; - # worker execution resources must remain outside this shared object. - self.linear_state_manager = req_manager.linear_state_manager - # super().write_to_shm() 会用 ForkingPickler 序列化本对象，torch 在 dump 时会把 - # CPU tensor 的 storage 原地迁到共享内存，使本进程大页 state cache 原本 - # pinned(cudaHostAlloc) 的内存退化为普通 shm mmap，之后 Triton kernel 携带该指针 - # 启动会报 "Pointer argument cannot be accessed from Triton (cpu tensor?)"。 - # 跨进程消费方并不使用 cpu 侧大页 state cache，序列化期间临时剔除以保住 pinned。 - big_page_buffers = self.big_page_buffers - self.big_page_buffers = None - try: - return super().write_to_shm(req_manager) - finally: - self.big_page_buffers = big_page_buffers - - def alloc_paged_kv_move_buffer(self, page_num, page_size) -> torch.Tensor: - kv_move_buffer = super().alloc_paged_kv_move_buffer(page_num, page_size) - self.att_state_page_helper = self._create_att_state_page_helper() - self.att_state_page_helper.assert_page_size() - return kv_move_buffer - - def _create_att_state_page_helper(self): - return Qwen3NextLinearAttPageHelper(self) - - def write_mem_to_page_kv_move_buffer( - self, - mem_indexes, - page_index: int, - dp_index: int, - mem_managers, - dp_world_size: int, - page_kind: str = "kv", - req_idx: int = None, - ): - if page_kind == "kv": - return super().write_mem_to_page_kv_move_buffer( - mem_indexes=mem_indexes, - page_index=page_index, - dp_index=dp_index, - mem_managers=mem_managers, - dp_world_size=dp_world_size, - page_kind=page_kind, - req_idx=req_idx, - ) - assert page_kind == "att_state", f"unknown page_kind={page_kind}" - assert req_idx is not None - helper = self.att_state_page_helper - dp_mems = helper.get_dp_mems(mem_managers, dp_index, dp_world_size) - helper.write_req_to_page(page_index=page_index, req_idx=req_idx, dp_mems=dp_mems) - return - - def read_page_kv_move_buffer_to_mem( - self, - mem_indexes, - page_index: int, - dp_index: int, - mem_managers, - dp_world_size: int, - page_kind: str = "kv", - req_idx: int = None, - ): - if page_kind == "kv": - return super().read_page_kv_move_buffer_to_mem( - mem_indexes=mem_indexes, - page_index=page_index, - dp_index=dp_index, - mem_managers=mem_managers, - dp_world_size=dp_world_size, - page_kind=page_kind, - req_idx=req_idx, - ) - assert page_kind == "att_state", f"unknown page_kind={page_kind}" - assert req_idx is not None - helper = self.att_state_page_helper - dp_mems = helper.get_dp_mems(mem_managers, dp_index, dp_world_size) - helper.read_page_to_req(page_index=page_index, req_idx=req_idx, dp_mems=dp_mems) - return - - -class _FP8StaticPerHeadQuantLinearAttMemOperator(LinearAttMemOperator): - def copy_kv_to_mem_manager(self, layer_index: int, mem_index: torch.Tensor, kv: torch.Tensor): - full_att_layer_index = self.linear_config.get_full_att_kv_layer_index(layer_index) - FP8StaticPerHeadQuantMemOperator.copy_kv_to_mem_manager(self, full_att_layer_index, mem_index, kv) - - -class _FP8StaticPerTensorQuantLinearAttMemOperator(LinearAttMemOperator): - def copy_kv_to_mem_manager(self, layer_index: int, mem_index: torch.Tensor, kv: torch.Tensor): - full_att_layer_index = self.linear_config.get_full_att_kv_layer_index(layer_index) - FP8StaticPerTensorQuantMemOperator.copy_kv_to_mem_manager(self, full_att_layer_index, mem_index, kv) - - -class FP8StaticPerHeadQuantQwen3NextMemManager(Qwen3NextMemManager, FP8StaticPerHeadQuantMemManager): - operator_class = _FP8StaticPerHeadQuantLinearAttMemOperator - - -class FP8StaticPerTensorQuantQwen3NextMemManager(Qwen3NextMemManager, FP8StaticPerTensorQuantMemManager): - operator_class = _FP8StaticPerTensorQuantLinearAttMemOperator - - -class Qwen3NextLinearAttPageHelper: - def __init__(self, mem_manager: "Qwen3NextMemManager"): - self.mem_manager = mem_manager - self.linear_config = mem_manager.linear_config - self.global_linear_k_heads = self.linear_config.global_linear_k_heads - self.global_linear_v_heads = self.linear_config.global_linear_v_heads - - self.global_q_dim = self.global_linear_k_heads * self.linear_config.head_linear_k_dim - self.global_k_dim = self.global_q_dim - self.global_v_heads = self.global_linear_v_heads - self.global_v_dim = self.global_v_heads * self.linear_config.head_linear_v_dim - # conv state follows mixed_qkv layout: [q, k, v], each as a flat channel block. - self.conv_shape = ( - self.linear_config.linear_layer_num, - self.global_q_dim + self.global_k_dim + self.global_v_dim, - self.linear_config.conv_kernel_size - 1, - ) - self.ssm_shape = ( - self.linear_config.linear_layer_num, - self.global_v_heads, - self.linear_config.head_linear_k_dim, - self.linear_config.head_linear_v_dim, - ) - - self.conv_nbytes = ( - self.conv_shape[0] * self.conv_shape[1] * self.conv_shape[2] * self.linear_config.conv_state_dtype.itemsize - ) - ssm_alignment = self.linear_config.ssm_state_dtype.itemsize - # 做一下字节对齐，防止切出来的不对齐，导致一些操作的性能下降。 - self.ssm_offset = ((self.conv_nbytes + ssm_alignment - 1) // ssm_alignment) * ssm_alignment - self.ssm_nbytes = ( - self.ssm_shape[0] - * self.ssm_shape[1] - * self.ssm_shape[2] - * self.ssm_shape[3] - * self.linear_config.ssm_state_dtype.itemsize - ) - self.state_nbytes = self.ssm_offset + self.ssm_nbytes - - def assert_page_size(self): - kv_move_buffer = self.mem_manager.kv_move_buffer - page_nbytes = kv_move_buffer[0].numel() * kv_move_buffer.element_size() - assert ( - page_nbytes >= self.state_nbytes - ), f"nixl kv move page bytes {page_nbytes} is smaller than global linear att state bytes {self.state_nbytes}" - return - - def get_dp_mems(self, mem_managers: List["Qwen3NextMemManager"], dp_index: int, dp_world_size: int): - dp_mems = mem_managers[(dp_index * dp_world_size) : ((dp_index + 1) * dp_world_size)] - assert len(dp_mems) == dp_world_size - for mem in dp_mems: - assert hasattr(mem, "linear_state_manager") - assert mem.linear_config.linear_layer_num == self.linear_config.linear_layer_num - assert mem.linear_config.conv_kernel_size == self.linear_config.conv_kernel_size - assert mem.linear_config.head_linear_k_dim == self.linear_config.head_linear_k_dim - assert mem.linear_config.head_linear_v_dim == self.linear_config.head_linear_v_dim - assert mem.linear_config.num_linear_k_heads == self.linear_config.num_linear_k_heads - assert mem.linear_config.num_linear_v_heads == self.linear_config.num_linear_v_heads - return dp_mems - - def view_page_to_linear_att_state(self, page_index: int): - page_bytes = self.mem_manager.kv_move_buffer[page_index].view(torch.uint8).reshape(-1) - conv_page = page_bytes[0 : self.conv_nbytes].view(self.linear_config.conv_state_dtype).view(self.conv_shape) - ssm_page = ( - page_bytes[self.ssm_offset : self.ssm_offset + self.ssm_nbytes] - .view(self.linear_config.ssm_state_dtype) - .view(self.ssm_shape) - ) - return conv_page, ssm_page - - def write_req_to_page( - self, - page_index: int, - req_idx: int, - dp_mems: List["Qwen3NextMemManager"], - ): - conv_page, ssm_page = self.view_page_to_linear_att_state(page_index) - for tp_index, mem in enumerate(dp_mems): - self._write_one_rank(mem, tp_index, req_idx, conv_page, ssm_page) - return - - def read_page_to_req( - self, - page_index: int, - req_idx: int, - dp_mems: List["Qwen3NextMemManager"], - ): - conv_page, ssm_page = self.view_page_to_linear_att_state(page_index) - for tp_index, mem in enumerate(dp_mems): - self._read_one_rank(mem, tp_index, req_idx, conv_page, ssm_page) - return - - def _write_one_rank( - self, - mem: "Qwen3NextMemManager", - tp_index: int, - req_idx: int, - conv_page: torch.Tensor, - ssm_page: torch.Tensor, - ): - conv_state, ssm_state = mem.linear_state_manager.get_accepted_state(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) - return - - def _copy_conv_state_to_page( - self, - conv_state: torch.Tensor, - conv_page: torch.Tensor, - mem: "Qwen3NextMemManager", - tp_index: int, - ): - local_q_heads = mem.linear_config.num_linear_k_heads - local_v_heads = mem.linear_config.num_linear_v_heads - head_k_dim = mem.linear_config.head_linear_k_dim - head_v_dim = mem.linear_config.head_linear_v_dim - - local_q_state = conv_state[:, 0 : local_q_heads * head_k_dim, :] - local_k_state = conv_state[:, local_q_heads * head_k_dim : 2 * local_q_heads * head_k_dim, :] - local_v_state = conv_state[:, 2 * local_q_heads * head_k_dim :, :] - global_q_page = conv_page[:, 0 : self.global_q_dim, :] - global_k_page = conv_page[:, self.global_q_dim : self.global_q_dim + self.global_k_dim, :] - global_v_page = conv_page[:, self.global_q_dim + self.global_k_dim :, :] - - qk_head_slice = self._get_head_slice( - tp_index, local_q_heads, self.global_linear_k_heads, mem.linear_config.tp_world_size, is_write=True - ) - if qk_head_slice is not None: - local_head_start, local_head_end, global_head_start, global_head_end = qk_head_slice - local_dim_start = local_head_start * head_k_dim - local_dim_end = local_head_end * head_k_dim - global_dim_start = global_head_start * head_k_dim - global_dim_end = global_head_end * head_k_dim - global_q_page[:, global_dim_start:global_dim_end, :].copy_( - local_q_state[:, local_dim_start:local_dim_end, :], non_blocking=True - ) - global_k_page[:, global_dim_start:global_dim_end, :].copy_( - local_k_state[:, local_dim_start:local_dim_end, :], non_blocking=True - ) - - v_head_slice = self._get_head_slice( - tp_index, local_v_heads, self.global_linear_v_heads, mem.linear_config.tp_world_size, is_write=True - ) - if v_head_slice is not None: - local_head_start, local_head_end, global_head_start, global_head_end = v_head_slice - local_dim_start = local_head_start * head_v_dim - local_dim_end = local_head_end * head_v_dim - global_dim_start = global_head_start * head_v_dim - global_dim_end = global_head_end * head_v_dim - global_v_page[:, global_dim_start:global_dim_end, :].copy_( - local_v_state[:, local_dim_start:local_dim_end, :], non_blocking=True - ) - return - - def _copy_ssm_state_to_page( - self, - ssm_state: torch.Tensor, - ssm_page: torch.Tensor, - mem: "Qwen3NextMemManager", - tp_index: int, - ): - head_slice = self._get_head_slice( - tp_index, - mem.linear_config.num_linear_v_heads, - self.global_linear_v_heads, - mem.linear_config.tp_world_size, - is_write=True, - ) - if head_slice is not None: - local_head_start, local_head_end, global_head_start, global_head_end = head_slice - ssm_page[:, global_head_start:global_head_end, :, :].copy_( - ssm_state[:, local_head_start:local_head_end, :, :], - non_blocking=True, - ) - return - - def _get_head_slice( - self, - tp_index: int, - local_heads: int, - global_heads: int, - tp_world_size: int, - is_write: bool, - ): - if local_heads == 0 or global_heads == 0: - return None - total_local_heads = local_heads * tp_world_size - repeat_count = max(1, total_local_heads // global_heads) - if is_write and repeat_count > 1 and tp_index % repeat_count != 0: - return None - unique_tp_index = tp_index // repeat_count - global_head_start = unique_tp_index * local_heads - global_head_end = min(global_head_start + local_heads, global_heads) - local_head_start = 0 - local_head_end = global_head_end - global_head_start - if local_head_end <= local_head_start: - return None - return local_head_start, local_head_end, global_head_start, global_head_end - - def _copy_page_to_conv_state( - self, - conv_page: torch.Tensor, - conv_state: torch.Tensor, - mem: "Qwen3NextMemManager", - tp_index: int, - ): - local_q_heads = mem.linear_config.num_linear_k_heads - local_v_heads = mem.linear_config.num_linear_v_heads - head_k_dim = mem.linear_config.head_linear_k_dim - head_v_dim = mem.linear_config.head_linear_v_dim - - local_q_state = conv_state[:, 0 : local_q_heads * head_k_dim, :] - local_k_state = conv_state[:, local_q_heads * head_k_dim : 2 * local_q_heads * head_k_dim, :] - local_v_state = conv_state[:, 2 * local_q_heads * head_k_dim :, :] - global_q_page = conv_page[:, 0 : self.global_q_dim, :] - global_k_page = conv_page[:, self.global_q_dim : self.global_q_dim + self.global_k_dim, :] - global_v_page = conv_page[:, self.global_q_dim + self.global_k_dim :, :] - - qk_head_slice = self._get_head_slice( - tp_index, local_q_heads, self.global_linear_k_heads, mem.linear_config.tp_world_size, is_write=False - ) - if qk_head_slice is not None: - local_head_start, local_head_end, global_head_start, global_head_end = qk_head_slice - local_dim_start = local_head_start * head_k_dim - local_dim_end = local_head_end * head_k_dim - global_dim_start = global_head_start * head_k_dim - global_dim_end = global_head_end * head_k_dim - local_q_state[:, local_dim_start:local_dim_end, :].copy_( - global_q_page[:, global_dim_start:global_dim_end, :], non_blocking=True - ) - local_k_state[:, local_dim_start:local_dim_end, :].copy_( - global_k_page[:, global_dim_start:global_dim_end, :], non_blocking=True - ) - - v_head_slice = self._get_head_slice( - tp_index, local_v_heads, self.global_linear_v_heads, mem.linear_config.tp_world_size, is_write=False - ) - if v_head_slice is not None: - local_head_start, local_head_end, global_head_start, global_head_end = v_head_slice - local_dim_start = local_head_start * head_v_dim - local_dim_end = local_head_end * head_v_dim - global_dim_start = global_head_start * head_v_dim - global_dim_end = global_head_end * head_v_dim - local_v_state[:, local_dim_start:local_dim_end, :].copy_( - global_v_page[:, global_dim_start:global_dim_end, :], non_blocking=True - ) - return - - def _copy_page_to_ssm_state( - self, - ssm_page: torch.Tensor, - ssm_state: torch.Tensor, - mem: "Qwen3NextMemManager", - tp_index: int, - ): - head_slice = self._get_head_slice( - tp_index, - mem.linear_config.num_linear_v_heads, - self.global_linear_v_heads, - mem.linear_config.tp_world_size, - is_write=False, - ) - if head_slice is not None: - local_head_start, local_head_end, global_head_start, global_head_end = head_slice - ssm_state[:, local_head_start:local_head_end, :, :].copy_( - ssm_page[:, global_head_start:global_head_end, :, :], - non_blocking=True, - ) - return - - def _read_one_rank( - self, - mem: "Qwen3NextMemManager", - tp_index: int, - req_idx: int, - conv_page: torch.Tensor, - ssm_page: torch.Tensor, - ): - conv_state, ssm_state = mem.linear_state_manager.prepare_restore(req_idx) - self._copy_page_to_conv_state(conv_page, conv_state, mem, tp_index) - self._copy_page_to_ssm_state(ssm_page, ssm_state, mem, tp_index) - return diff --git a/lightllm/common/linear_state.py b/lightllm/common/linear_state.py index bfa913d4fd..d39fb40813 100644 --- a/lightllm/common/linear_state.py +++ b/lightllm/common/linear_state.py @@ -10,7 +10,6 @@ import torch -from lightllm.common.state_cache_manager import LayerCache, LinearAttCacheConfig if TYPE_CHECKING: from lightllm.common.basemodel.triton_kernel.linear_att.replayssm import ReplaySSMBackend @@ -40,76 +39,77 @@ class LinearStateTensors(NamedTuple): class LinearStateManager: - """Own request state and its lifecycle, without dispatching attention. - - PD currently shares this object through CUDA IPC as part of the memory - manager. Keep its fields limited to tensors, cache operators, and pickleable - static configuration. Graphs, streams, callbacks, batch contexts, and pinned - CPU checkpoint pools belong to their existing worker/runner owners. Adding - such resources requires an explicit transport view rather than serializing - them implicitly with this object. - """ - - def __init__(self, max_request_num, linear_config: LinearAttCacheConfig, args, *, kda=False, device="cuda"): - self.linear_config = linear_config - self.mtp_step = args.mtp_step - self.ssm_slots_per_req = self.mtp_step + 1 if args.ssm_state_mode == "native" else 1 - self.accepted_state_index = ( - torch.zeros(max_request_num + 1, dtype=torch.int32, device=device) if self.mtp_step > 0 else None - ) - self.conv_state = LayerCache( - size=max_request_num + 1, - dtype=linear_config.conv_state_dtype, - shape=linear_config.get_mtp_conv_state_shape(mtp_step=self.mtp_step), - layer_num=linear_config.linear_layer_num, + """Shared Conv/SSM lifecycle for GDN and KDA attention resources.""" + + def __init__( + self, + config, + request_capacity, + mtp_step=0, + device="cuda", + *, + ssm_state_mode="native", + ssm_replay_config=None, + kda=False, + gate_lower_bound=-5.0, + ): + self.config = config + self.request_capacity = request_capacity + self.mtp_step = mtp_step + self.ssm_slots_per_req = mtp_step + 1 if ssm_state_mode == "native" else 1 + self.mtp_state_index = torch.zeros(request_capacity, dtype=torch.int32, device=device) if mtp_step else None + self._logical_to_storage_layer = {layer: index for index, layer in enumerate(config.logical_layer_ids)} + layers = len(config.logical_layer_ids) + self.conv_state = torch.zeros( + (layers, request_capacity, config.conv_channels, config.conv_width + mtp_step), + dtype=config.conv_dtype, device=device, ) - self.ssm_state = LayerCache( - size=(max_request_num + 1) * self.ssm_slots_per_req, - dtype=linear_config.ssm_state_dtype, - shape=linear_config.get_ssm_state_shape(), - layer_num=linear_config.linear_layer_num, + self.ssm_state = torch.zeros( + (layers, request_capacity * self.ssm_slots_per_req, config.num_v_heads, config.key_dim, config.value_dim), + dtype=config.ssm_dtype, device=device, ) + ssm_replay_config = ssm_replay_config or {} self.ssm_backend: Optional[Union["ReplaySSMBackend", "CompactSSMBackend"]] = None - if args.ssm_state_mode == "replay": + if ssm_state_mode == "replay": from lightllm.common.basemodel.triton_kernel.linear_att import replayssm self.ssm_backend = replayssm.ReplaySSMBackend( - self.ssm_state.buffer, - args.ssm_replay_config.get("cache_len", 16), + self.ssm_state, + ssm_replay_config.get("cache_len", 16), self.mtp_step + 1, - linear_config.conv_state_dtype, - num_key_heads=linear_config.num_linear_k_heads, + config.conv_dtype, + num_key_heads=config.num_k_heads, kda=kda, - gate_lower_bound=linear_config.gate_lower_bound, - compute_mode=args.ssm_replay_config.get("compute_mode", "inline"), - fold_method=args.ssm_replay_config.get("fold_method", "raw"), + gate_lower_bound=gate_lower_bound, + compute_mode=ssm_replay_config.get("compute_mode", "inline"), + fold_method=ssm_replay_config.get("fold_method", "raw"), ) - elif args.ssm_state_mode == "compact": + elif ssm_state_mode == "compact": from lightllm.common.basemodel.triton_kernel.linear_att import compact_ssm self.ssm_backend = compact_ssm.CompactSSMBackend( - self.ssm_state.buffer, + self.ssm_state, self.mtp_step + 1, - linear_config.conv_state_dtype, - num_key_heads=linear_config.num_linear_k_heads, + config.conv_dtype, + num_key_heads=config.num_k_heads, kda=kda, - gate_lower_bound=linear_config.gate_lower_bound, + gate_lower_bound=gate_lower_bound, ) def _reset_tracking(self, req_idx): if self.ssm_backend is not None: self.ssm_backend.clear_history(req_idx) - if self.accepted_state_index is not None: - self.accepted_state_index[req_idx].zero_() + if self.mtp_state_index is not None: + self.mtp_state_index[req_idx].zero_() def reset(self, req_idx): """Initialize a newly allocated/reused request, including speculative slots.""" self._reset_tracking(req_idx) - self.conv_state.buffer[:, req_idx].zero_() + self.conv_state[:, req_idx].zero_() start = req_idx * self.ssm_slots_per_req - self.ssm_state.buffer[:, start : start + self.ssm_slots_per_req].zero_() + self.ssm_state[:, start : start + self.ssm_slots_per_req].zero_() def prepare_prefill(self, reqs): """Merge accepted history and return convolution and SSM buffer indexes.""" @@ -127,8 +127,8 @@ def prepare_decode(self, reqs, cu_seqlens=None, *, single_token=False): def prepare_mtp_decode(self, reqs, cu_seqlens): """Return Conv/SSM indexes and prior accepted token counts; prepare replay history.""" - accepted_token_counts = self.accepted_state_index[reqs] + 1 - accepted_token_counts = torch.where(reqs == self.accepted_state_index.numel() - 1, 1, accepted_token_counts) + accepted_token_counts = self.mtp_state_index[reqs] + 1 + accepted_token_counts = torch.where(reqs == self.mtp_state_index.numel() - 1, 1, accepted_token_counts) if self.ssm_backend is not None: ssm_indexes = reqs else: @@ -166,7 +166,7 @@ def commit_accepted_state( if isinstance(self.ssm_backend, replayssm.ReplaySSMBackend) and self.ssm_backend.verify_width > 1: replay_cursors = self.ssm_backend.cursors linear_att_mtp_state_index_update( - req_to_mtp_state_index=self.accepted_state_index, + req_to_mtp_state_index=self.mtp_state_index, b_req_mtp_start_loc=b_req_mtp_start_loc, b_req_idx=b_req_idx, b_mtp_index=b_mtp_index, @@ -178,7 +178,7 @@ def commit_accepted_state( return reqs = b_req_idx[b_req_mtp_start_loc.long()] if linear_state_batches is None or isinstance(self.ssm_backend, replayssm.ReplaySSMBackend): - self.ssm_backend.accept_updates(reqs, self.accepted_state_index) + self.ssm_backend.accept_updates(reqs, self.mtp_state_index) return assert sum(size for size, _ in linear_state_batches) == reqs.numel() start = 0 @@ -186,7 +186,7 @@ def commit_accepted_state( if size: self.ssm_backend.accept_updates( reqs[start : start + size], - self.accepted_state_index, + self.mtp_state_index, run_config=commit._run_config if commit is not None else None, ) start += size @@ -200,18 +200,18 @@ def get_accepted_state(self, req_idx): cross-process callers must use their existing completion events before recycling sources or destinations. No CPU transfer or sync is implicit. """ - width = self.linear_config.get_conv_state_shape()[-1] - conv = self.conv_state.buffer[:, req_idx, ..., :width] - if self.accepted_state_index is not None: - offsets = torch.arange(width, device=conv.device) + self.accepted_state_index[req_idx] - conv = self.conv_state.buffer[:, req_idx].index_select(-1, offsets) + width = self.config.conv_width + conv = self.conv_state[:, req_idx, ..., :width] + if self.mtp_state_index is not None: + offsets = torch.arange(width, device=conv.device) + self.mtp_state_index[req_idx] + conv = self.conv_state[:, req_idx].index_select(-1, offsets) if self.ssm_backend is not None: ssm = self.ssm_backend.get_accepted_state(req_idx) elif self.ssm_slots_per_req > 1: - index = req_idx * self.ssm_slots_per_req + self.accepted_state_index[req_idx] - ssm = self.ssm_state.buffer.index_select(1, index.long().view(1))[:, 0] + index = req_idx * self.ssm_slots_per_req + self.mtp_state_index[req_idx] + ssm = self.ssm_state.index_select(1, index.long().view(1))[:, 0] else: - ssm = self.ssm_state.buffer[:, req_idx] + ssm = self.ssm_state[:, req_idx] return LinearStateTensors(conv, ssm) def prepare_restore(self, req_idx): @@ -222,10 +222,10 @@ def prepare_restore(self, req_idx): The same stream/completion rules as ``get_accepted_state`` apply. """ self._reset_tracking(req_idx) - width = self.linear_config.get_conv_state_shape()[-1] + width = self.config.conv_width return LinearStateTensors( - self.conv_state.buffer[:, req_idx, ..., :width], - self.ssm_state.buffer[:, req_idx * self.ssm_slots_per_req], + self.conv_state[:, req_idx, ..., :width], + self.ssm_state[:, req_idx * self.ssm_slots_per_req], ) def restore(self, req_idx, conv, ssm): @@ -233,34 +233,14 @@ def restore(self, req_idx, conv, ssm): target.conv.copy_(conv, non_blocking=True) target.ssm.copy_(ssm, non_blocking=True) - def save_checkpoints(self, b_req_idx, req_indexes, buffer_indexes, destination): - """Copy current accepted prefixes into caller-allocated checkpoint slots.""" - if self.ssm_backend is None: - from lightllm.common.basemodel.triton_kernel.linear_att_copy import copy_linear_att_state_to_kv_buffer + def get_storage_layer_index(self, logical_layer_id): + return self._logical_to_storage_layer[logical_layer_id] - indexes = torch.tensor(buffer_indexes, dtype=torch.int32, device="cpu").to( - b_req_idx.device, non_blocking=True - ) - copy_linear_att_state_to_kv_buffer( - b_req_idx=b_req_idx, - big_page_buffer_ids=indexes, - gpu_conv_state=self.conv_state.buffer, - gpu_ssm_state=self.ssm_state.buffer, - cpu_kv_conv_state=destination.conv_state_cache.buffer, - cpu_kv_ssm_state=destination.ssm_state_cache.buffer, - mtp_step=self.ssm_slots_per_req - 1, - accepted_state_indices=self.accepted_state_index, - ) - else: - for req_idx, buffer_idx in zip(req_indexes, buffer_indexes): - if buffer_idx != -1: - self.copy_accepted_state_to(req_idx, *destination.get_state_cache(buffer_idx)) - - def copy_accepted_state_to(self, req_idx, conv, ssm): - accepted_state = self.get_accepted_state(req_idx) - conv.copy_(accepted_state.conv, non_blocking=True) - ssm.copy_(accepted_state.ssm, non_blocking=True) - - def get_layer_buffers(self, layer_idx_in_all): - linear_layer_idx = self.linear_config.get_linear_layer_index(layer_idx_in_all) - return self.conv_state.buffer[linear_layer_idx], self.ssm_state.buffer[linear_layer_idx] + +def get_linear_state(infer_state, state_class): + (state,) = ( + state + for state in infer_state.att_state_manager.get_bounded_attention_states(infer_state.att_state_role) + if isinstance(state, state_class) + ) + return state diff --git a/lightllm/common/req_manager/glm5_next.py b/lightllm/common/req_manager/glm5_next.py deleted file mode 100644 index 5814404b65..0000000000 --- a/lightllm/common/req_manager/glm5_next.py +++ /dev/null @@ -1,36 +0,0 @@ -from lightllm.common.state_cache_manager import LayerCache -from lightllm.utils.envs_utils import get_env_start_args - -from .linear_att import ReqManagerForMamba - - -class Glm5NextReqManager(ReqManagerForMamba): - """KDA runtime state and the NSA indexer's incomplete four-token pool.""" - - kda = True - - def __init__(self, max_request_num, max_sequence_length, mem_manager, linear_config): - # Both checkpoint sizes are multiples of this hash page. A restored - # prefix therefore has no incomplete K-pool to serialize or replay. - assert ( - get_env_start_args().linear_att_hash_page_size % linear_config.index_kpool == 0 - ), "GLM K-pool requires cache pages aligned to index_kpool" - super().__init__(max_request_num, max_sequence_length, mem_manager, linear_config) - self.req_to_indexer_tail = LayerCache( - size=max_request_num + 1, - dtype=linear_config.full_att_dtype, - shape=(linear_config.index_kpool + self.mtp_step, 2 * linear_config.index_head_dim), - layer_num=linear_config.get_full_att_kv_layer_num_with_draft_model(), - device="cuda", - ) - - def get_indexer_tail_buffer(self, layer_index): - return self.req_to_indexer_tail.buffer[self.linear_config.get_full_att_kv_layer_index(layer_index)] - - def init_hybrid_attention_state(self, req): - super().init_hybrid_attention_state(req) - self.req_to_indexer_tail.buffer[:, req.req_idx].zero_() - - def restore_state(self, req, state_cache_manager, buffer_idx): - super().restore_state(req, state_cache_manager, buffer_idx) - self.req_to_indexer_tail.buffer[:, req.req_idx].zero_() diff --git a/lightllm/common/req_manager/hybrid_base.py b/lightllm/common/req_manager/hybrid_base.py deleted file mode 100644 index a1d178edec..0000000000 --- a/lightllm/common/req_manager/hybrid_base.py +++ /dev/null @@ -1,72 +0,0 @@ -from abc import ABC, abstractmethod -from typing import TYPE_CHECKING, List - -import torch - -from .base import ReqManager - - -if TYPE_CHECKING: - from lightllm.server.router.model_infer.infer_batch import InferReq - - -class HybridAttentionReqManager(ReqManager, ABC): - """混合 attention 的请求运行态与大小页 checkpoint 管理接口。 - - 大小页沿同一虚拟 token 索引空间匹配前缀，full attention KV 保持 token 粒度存储。 - linear/sliding-window 状态在大页边界及请求可缓存尾部的小页边界保存 checkpoint， - 缓存命中后，再将相应 checkpoint 恢复到请求运行态。 - - 请求的 GPU 计算状态由具体实现管理；small_page_buffers 持有 CPU 小页快照； - big_page_buffers 引用 mem_manager 的 CPU 大页快照，与 full KV 容量一起创建、调整。 - 大小页不包含额外的 GPU 运行态，命中后恢复到请求状态，不重新分配 checkpoint 池。 - 公共缓存流程负责 checkpoint 槽位分配、边界、匹配与淘汰；本接口负责运行态与 checkpoint 保存/恢复。 - """ - - def __init__(self, max_request_num, max_sequence_length, mem_manager): - super().__init__(max_request_num, max_sequence_length, mem_manager) - self.small_page_buffers = None - - @property - def big_page_buffers(self): - return self.mem_manager.big_page_buffers - - @abstractmethod - def create_small_page_cache_manager(self, size: int): - """创建并持有 CPU 小页池，返回同一池供 radix 使用；size 是槽位数，不是 token 数。""" - - @abstractmethod - def init_hybrid_attention_state(self, req: "InferReq"): - """无前缀缓存命中时，初始化已分配请求槽位的 GPU 运行态。""" - - def restore_big_page_state(self, big_page_buffer_idx: int, req: "InferReq"): - """将指定大页槽位的 CPU checkpoint 恢复到请求 GPU 运行态。""" - self.restore_state(req, self.big_page_buffers, big_page_buffer_idx) - - def restore_small_page_state(self, req: "InferReq"): - """将 req.shared_kv_node 对应的小页 checkpoint 恢复到请求 GPU 运行态。""" - self.restore_state(req, self.small_page_buffers, req.shared_kv_node.small_page_buffer_idx) - - @abstractmethod - def restore_state(self, req: "InferReq", state_cache_manager, buffer_idx: int): - """CPU checkpoint → 请求 GPU 运行态；大小页共用，不负责前缀匹配或 full KV 索引恢复。""" - - def save_big_page_states(self, b_req_idx: torch.Tensor, req_indexes: List[int], buffer_indexes: List[int]): - """批量保存请求 GPU 运行态到已分配的大页槽位，buffer_indexes 中的 -1 表示跳过。 - - b_req_idx 与 req_indexes 分别为同一批请求的 GPU 索引张量和 CPU 索引列表。 - 默认逐请求保存，模型可覆盖为批量拷贝算子。 - """ - for req_idx, buffer_idx in zip(req_indexes, buffer_indexes): - if buffer_idx != -1: - self.save_state(req_idx, buffer_idx, self.big_page_buffers) - - @abstractmethod - def save_state(self, req_idx: int, buffer_idx: int, state_cache_manager): - """请求 GPU 运行态 → 指定 CPU checkpoint 槽位；大小页共用，调用方负责分配槽位。""" - - def update_mtp_state( - self, b_req_mtp_start_loc, b_req_idx, b_mtp_index, accepted_index, verify_width, linear_state_batches=None - ): - """接受推测 token 后更新运行态位置；需要调整状态索引的模型覆写。""" - return diff --git a/lightllm/common/req_manager/linear_att.py b/lightllm/common/req_manager/linear_att.py deleted file mode 100644 index 41f22dbfe4..0000000000 --- a/lightllm/common/req_manager/linear_att.py +++ /dev/null @@ -1,55 +0,0 @@ -from typing import TYPE_CHECKING, List - -import torch - -from lightllm.common.linear_state import LinearStateManager -from lightllm.common.state_cache_manager import LinearAttCacheConfig, LinearAttCacheManager -from lightllm.utils.envs_utils import get_env_start_args - -from .hybrid_base import HybridAttentionReqManager - - -if TYPE_CHECKING: - from lightllm.server.router.model_infer.infer_batch import InferReq - - -class ReqManagerForMamba(HybridAttentionReqManager): - kda = False - - def __init__(self, max_request_num, max_sequence_length, mem_manager, linear_config: LinearAttCacheConfig): - super().__init__(max_request_num, max_sequence_length, mem_manager) - args = get_env_start_args() - self.mtp_step = args.mtp_step - # Prefill kernels read the canonical convolution window. Mixed prefill - # and MTP decode still require a separate execution-layout change. - if self.mtp_step > 0: - assert args.enable_prefill_decode_mixed is False - self.big_page_token_num = args.linear_att_page_block_num * args.linear_att_hash_page_size - self.linear_config = linear_config - self.linear_state_manager = LinearStateManager(max_request_num, linear_config, args, kda=self.kda) - - def init_hybrid_attention_state(self, req: "InferReq"): - self.linear_state_manager.reset(req.req_idx) - - def create_small_page_cache_manager(self, size: int): - self.small_page_buffers = LinearAttCacheManager(size=size, linear_config=self.linear_config) - return self.small_page_buffers - - def save_big_page_states(self, b_req_idx: torch.Tensor, req_indexes: List[int], buffer_indexes: List[int]): - self.linear_state_manager.save_checkpoints(b_req_idx, req_indexes, buffer_indexes, self.big_page_buffers) - - def save_state(self, req_idx: int, buffer_idx: int, state_cache_manager: LinearAttCacheManager): - self.linear_state_manager.copy_accepted_state_to(req_idx, *state_cache_manager.get_state_cache(buffer_idx)) - - def get_mamba_cache(self, layer_idx_in_all: int): - return self.linear_state_manager.get_layer_buffers(layer_idx_in_all) - - def update_mtp_state( - self, b_req_mtp_start_loc, b_req_idx, b_mtp_index, accepted_index, verify_width, linear_state_batches=None - ): - self.linear_state_manager.commit_accepted_state( - b_req_mtp_start_loc, b_req_idx, b_mtp_index, accepted_index, verify_width, linear_state_batches - ) - - def restore_state(self, req: "InferReq", state_cache_manager: LinearAttCacheManager, buffer_idx: int): - self.linear_state_manager.restore(req.req_idx, *state_cache_manager.get_state_cache(buffer_idx)) diff --git a/lightllm/common/state_cache_manager/linear_att.py b/lightllm/common/state_cache_manager/linear_att.py deleted file mode 100644 index ec6cae7666..0000000000 --- a/lightllm/common/state_cache_manager/linear_att.py +++ /dev/null @@ -1,205 +0,0 @@ -import dataclasses - -import torch -import triton - -from lightllm.utils.envs_utils import get_added_mtp_kv_layer_num, get_env_start_args -from lightllm.utils.log_utils import init_logger -from lightllm.utils.torch_dtype_utils import get_torch_dtype - -from .base import StateCacheManager -from .layer_cache import LayerCache - - -logger = init_logger(__name__) - - -@dataclasses.dataclass -class LinearAttCacheConfig: - tp_world_size: int - # full att 的参数 - full_att_all_num_kv_heads: int - full_att_dtype: torch.dtype - full_att_num_kv_heads: int # 这个是 tp 后的head头数量 - full_att_head_dim: int - - # linear att 的参数 - global_linear_k_heads: int - global_linear_v_heads: int - num_linear_k_heads: int - num_linear_v_heads: int - head_linear_k_dim: int - head_linear_v_dim: int - conv_kernel_size: int - linear_layer_num: int - conv_state_dtype: torch.dtype - ssm_state_dtype: torch.dtype - full_attention_interval: int - all_layer_num: int # 包括 linear att 和 full att 的层加起来的层数 - draft_full_att_kv_layer_num: int = 0 - # Static KDA recurrence parameter; unused by GDN. - gate_lower_bound: float = -5.0 - - def get_conv_dim(self): - # 第一项对应q的参数，第二项对应k的参数，第三项对应v的参数 - # 由于 k_dim = q_dim, k_heads = q_heads, 所以第一项和第二项的计算 - # 形式相同，但是实际内在含义是不同的。 - return ( - self.head_linear_k_dim * self.num_linear_k_heads - + self.head_linear_k_dim * self.num_linear_k_heads - + self.head_linear_v_dim * self.num_linear_v_heads - ) - - def get_main_model_full_att_layer_num(self): - full_att_layer_num = self.all_layer_num - self.linear_layer_num - assert full_att_layer_num == self.all_layer_num // self.full_attention_interval - return full_att_layer_num - - def get_full_att_kv_layer_num_with_draft_model(self): - return self.get_main_model_full_att_layer_num() + self.draft_full_att_kv_layer_num - - def get_full_att_kv_layer_index(self, layer_index: int) -> int: - """Map a global target/draft layer index to the packed full-attention cache.""" - - layer_index = int(layer_index) - if layer_index >= self.all_layer_num: - kv_layer_index = layer_index - self.linear_layer_num - else: - kv_layer_index = layer_index // self.full_attention_interval - assert 0 <= kv_layer_index < self.get_full_att_kv_layer_num_with_draft_model(), ( - f"layer {layer_index} maps outside the packed full-attention KV cache: " - f"slot={kv_layer_index}, slots={self.get_full_att_kv_layer_num_with_draft_model()}" - ) - return kv_layer_index - - def get_linear_layer_index(self, layer_index: int) -> int: - """Map a target linear-attention layer to its packed recurrent-state row.""" - assert 0 <= layer_index < self.all_layer_num, f"invalid transformer layer index {layer_index}" - assert (layer_index + 1) % self.full_attention_interval != 0, f"layer {layer_index} uses full attention" - linear_index = layer_index - layer_index // self.full_attention_interval - assert linear_index < self.linear_layer_num, f"layer {layer_index} maps outside the recurrent-state pool" - return linear_index - - def get_conv_state_shape(self): - # Base committed sliding-window state, without speculative MTP tail. - return (self.get_conv_dim(), self.conv_kernel_size - 1) - - def get_mtp_conv_state_shape(self, mtp_step: int): - # Working state with room for S speculative tokens before acceptance. - return (self.get_conv_dim(), (self.conv_kernel_size - 1) + mtp_step) - - def get_ssm_state_shape(self): - return (self.num_linear_v_heads, self.head_linear_k_dim, self.head_linear_v_dim) - - def get_conv_state_bytes_per_layer(self): - return self.get_conv_dim() * (self.conv_kernel_size - 1) * self.conv_state_dtype.itemsize - - def get_ssm_state_bytes_per_layer(self): - return self.num_linear_v_heads * self.head_linear_k_dim * self.head_linear_v_dim * self.ssm_state_dtype.itemsize - - def get_cpu_cache_big_page_bytes(self): - a = self.get_cpu_cache_full_att_bytes() - b = self.get_cpu_cache_conv_bytes() - c = self.get_cpu_cache_ssm_bytes() - - return triton.cdiv(a + b + c, 16) * 16 - - def get_cpu_cache_full_att_bytes(self): - big_page_token_num = ( - get_env_start_args().linear_att_page_block_num * get_env_start_args().linear_att_hash_page_size - ) - assert big_page_token_num == get_env_start_args().cpu_cache_token_page_size - full_att_bytes = 2 * self.full_att_all_num_kv_heads * self.full_att_head_dim * self.full_att_dtype.itemsize - a = full_att_bytes * self.get_full_att_kv_layer_num_with_draft_model() * big_page_token_num - return a - - def get_cpu_cache_conv_bytes(self): - b = self.get_conv_state_bytes_per_layer() * self.linear_layer_num * self.tp_world_size - return b - - def get_cpu_cache_ssm_bytes(self): - c = self.get_ssm_state_bytes_per_layer() * self.linear_layer_num * self.tp_world_size - return c - - @staticmethod - def load_from_args() -> "LinearAttCacheConfig": - args = get_env_start_args() - model_path = args.model_dir - from transformers.configuration_utils import PretrainedConfig - - model_cfg, _ = PretrainedConfig.get_config_dict(model_path) - model_type = model_cfg["model_type"] - assert model_type in ["qwen3_5", "qwen3_5_moe", "qwen3_5_text", "qwen3_5_moe_text"] - llm_config = model_cfg - try: - llm_config = llm_config["text_config"] - except: - pass - - n_layer = llm_config["num_hidden_layers"] - - tp_world_size = get_env_start_args().tp // get_env_start_args().dp - full_att_dtype = ( - torch.uint8 if args.llm_kv_type in {"fp8kv_sph", "fp8kv_spt"} else get_torch_dtype(args.data_type) - ) - return LinearAttCacheConfig( - tp_world_size=tp_world_size, - full_att_all_num_kv_heads=llm_config["num_key_value_heads"], - full_att_dtype=full_att_dtype, - full_att_num_kv_heads=max(1, llm_config["num_key_value_heads"] // tp_world_size), - full_att_head_dim=llm_config["head_dim"], - global_linear_k_heads=llm_config["linear_num_key_heads"], - global_linear_v_heads=llm_config["linear_num_value_heads"], - num_linear_k_heads=max(1, llm_config["linear_num_key_heads"] // tp_world_size), - num_linear_v_heads=max(1, llm_config["linear_num_value_heads"] // tp_world_size), - head_linear_k_dim=llm_config["linear_key_head_dim"], - head_linear_v_dim=llm_config["linear_value_head_dim"], - conv_kernel_size=llm_config["linear_conv_kernel_dim"], - linear_layer_num=n_layer - (n_layer // llm_config["full_attention_interval"]), - conv_state_dtype=get_torch_dtype(args.data_type), - ssm_state_dtype=get_torch_dtype(args.linear_att_ssm_data_type), - full_attention_interval=llm_config["full_attention_interval"], - all_layer_num=n_layer, - draft_full_att_kv_layer_num=get_added_mtp_kv_layer_num(), - ) - - -class LinearAttCacheManager(StateCacheManager): - """CPU pinned conv/SSM checkpoint，两个 buffer 均保持 size-first 布局。""" - - def __init__( - self, - size: int, - linear_config: LinearAttCacheConfig, - keep_num: int = 0, # 用于记录需要保留的缓存数量，用于支持含有 linear_att 的如qwen3.5 模型的cpu cache的碎页处理。 - ): - super().__init__(size, keep_num) - self.linear_config = linear_config - # init the layer cache - self.conv_state_cache = LayerCache( - size=self.size, - dtype=self.linear_config.conv_state_dtype, - shape=self.linear_config.get_conv_state_shape(), - layer_num=self.linear_config.linear_layer_num, - device="cpu", - size_first=True, - ) - self.ssm_state_cache = LayerCache( - size=self.size, - dtype=self.linear_config.ssm_state_dtype, - shape=self.linear_config.get_ssm_state_shape(), - layer_num=self.linear_config.linear_layer_num, - device="cpu", - size_first=True, - ) - self.clear_to_init_state() - return - - def get_state_cache(self, buffer_idx: int): - return self.conv_state_cache.buffer[buffer_idx, ...], self.ssm_state_cache.buffer[buffer_idx, ...] - - def clear_to_init_state(self): - self.conv_state_cache.buffer.zero_() - self.ssm_state_cache.buffer.zero_() - super().clear_to_init_state() - return diff --git a/lightllm/models/glm5_next/config.py b/lightllm/models/glm5_next/config.py index 69abc162ba..f19f9028d3 100644 --- a/lightllm/models/glm5_next/config.py +++ b/lightllm/models/glm5_next/config.py @@ -17,6 +17,7 @@ class Glm5NextResourceConfig(Qwen3NextResourceConfig): """Replicated MLA/index KV plus TP-sharded KDA checkpoints.""" # Indexer 池化步长，单位为原始 token；同时决定 K-pool tail 的基准 ring 长度和恢复边界。 + gate_lower_bound: float = -5.0 index_kpool: int = 4 # 单个 indexer key 的特征维度；原始 key/gate tail 的 channel 宽度为该值的两倍。 index_head_dim: int = 128 diff --git a/lightllm/models/glm5_next/model.py b/lightllm/models/glm5_next/model.py index 5b5235c76e..2b2977bcb9 100644 --- a/lightllm/models/glm5_next/model.py +++ b/lightllm/models/glm5_next/model.py @@ -79,6 +79,9 @@ def create_bounded_att_states( request_capacity=request_capacity, mtp_step=args.mtp_step, device=device, + ssm_state_mode=args.ssm_state_mode, + ssm_replay_config=args.ssm_replay_config, + gate_lower_bound=cfg.gate_lower_bound, ) ) full_layers = tuple(model_logical_layer_ids[index] for index in cfg.get_full_att_layer_indices()) diff --git a/lightllm/models/qwen3next/model.py b/lightllm/models/qwen3next/model.py index f1a3a65748..20fbeb4749 100644 --- a/lightllm/models/qwen3next/model.py +++ b/lightllm/models/qwen3next/model.py @@ -81,6 +81,8 @@ def create_bounded_att_states( request_capacity=request_capacity, mtp_step=args.mtp_step, device=device, + ssm_state_mode=args.ssm_state_mode, + ssm_replay_config=args.ssm_replay_config, ) return [state] diff --git a/lightllm/server/router/model_infer/infer_batch.py b/lightllm/server/router/model_infer/infer_batch.py index f5d712a3fb..640c049aa3 100644 --- a/lightllm/server/router/model_infer/infer_batch.py +++ b/lightllm/server/router/model_infer/infer_batch.py @@ -1,9 +1,6 @@ import enum import collections import pickle -from dataclasses import dataclass, field -from typing import TYPE_CHECKING, List, Dict, Tuple, Optional, Callable, Any, Union - import torch import numpy as np from sortedcontainers import SortedDict diff --git a/lightllm/server/router/model_infer/mtp_speculative/utils.py b/lightllm/server/router/model_infer/mtp_speculative/utils.py index 7b68d06179..98fe1ea809 100644 --- a/lightllm/server/router/model_infer/mtp_speculative/utils.py +++ b/lightllm/server/router/model_infer/mtp_speculative/utils.py @@ -5,6 +5,8 @@ import torch +from lightllm.common.linear_state import LinearStateManager + from lightllm.common.basemodel.triton_kernel.mtp_utils import ( mtp_scatter_next_token_ids, mtp_verify, @@ -25,7 +27,6 @@ def verify_mtp_tokens( linear_state_batches=None, ) -> Tuple[torch.Tensor, torch.Tensor]: """Verify target tokens and update recurrent MTP state when required.""" - # 保留原调用接口中的步序参数；当前状态提交直接使用验证返回的接受长度。 accept_lengths, accepted_index = mtp_verify( req_to_next_token_ids=backend.model.req_manager.req_sampling_params_manager.req_to_next_token_ids, @@ -33,15 +34,18 @@ def verify_mtp_tokens( new_next_token_ids=next_token_ids, b_req_idx=b_req_idx, ) - if backend.is_hybrid_att_model: - backend.model.req_manager.update_mtp_state( - b_req_mtp_start_loc=b_req_mtp_start_loc, - b_req_idx=b_req_idx, - b_mtp_index=b_mtp_index, - accepted_index=accepted_index, - verify_width=backend.max_draft_step + 1, - linear_state_batches=linear_state_batches, - ) + for state in backend.model.att_state_manager.get_bounded_attention_states("main"): + if isinstance(state, LinearStateManager): + state.commit_accepted_state( + b_req_mtp_start_loc=b_req_mtp_start_loc, + b_req_idx=b_req_idx, + b_mtp_index=b_mtp_index, + accepted_index=accepted_index, + verify_width=backend.max_draft_step + 1, + linear_state_batches=linear_state_batches, + ) + else: + state.commit_bounded_att_state(b_req_idx[b_req_mtp_start_loc], accept_lengths) return accept_lengths, accepted_index diff --git a/unit_tests/common/att_resources/model_helpers.py b/unit_tests/common/att_resources/model_helpers.py index 10c52c8cea..86e62cb671 100644 --- a/unit_tests/common/att_resources/model_helpers.py +++ b/unit_tests/common/att_resources/model_helpers.py @@ -253,6 +253,8 @@ def build_hybrid_resources( small_page_capacity: int, draft_full_state: FullAttentionState | None = None, tp_rank: int = 0, + ssm_state_mode: str = "native", + ssm_replay_config: dict | None = None, ) -> AttentionStateManager: """测试按既定几何参数组装通用管理器，不依赖旧模型管理器的构造/绑定流程。""" from lightllm.common.attention_state import AttentionStateManager @@ -282,6 +284,9 @@ def build_hybrid_resources( request_capacity=request_capacity, mtp_step=mtp_step, device="cuda", + ssm_state_mode=ssm_state_mode, + ssm_replay_config=ssm_replay_config, + gate_lower_bound=config.gate_lower_bound if glm else -5.0, ) bounded_state.initialize_page_pools( big_page_capacity=big_page_capacity, diff --git a/unit_tests/common/att_resources/test_hybrid_layer_mapping.py b/unit_tests/common/att_resources/test_hybrid_layer_mapping.py index 15d29bd85d..f93a87c1e9 100644 --- a/unit_tests/common/att_resources/test_hybrid_layer_mapping.py +++ b/unit_tests/common/att_resources/test_hybrid_layer_mapping.py @@ -9,6 +9,7 @@ from lightllm.server.core.objs.start_args_type import StartArgs +@pytest.mark.parametrize("mode", ["native", "compact", "replay"]) @pytest.mark.parametrize("draft_layers", [0, 1, 3]) @pytest.mark.parametrize( "layer_types,expected_indices,linear_layers", @@ -41,6 +42,7 @@ def test_hybrid_layer_mapping_and_mtp_states( expected_indices: tuple[int, ...], linear_layers: int, draft_layers: int, + mode: str, ) -> None: model_config = dict( num_hidden_layers=len(layer_types), @@ -56,7 +58,13 @@ def test_hybrid_layer_mapping_and_mtp_states( full_attention_interval=4, layer_types=layer_types, ) - args = StartArgs(data_type="bfloat16", linear_att_ssm_data_type="float32") + args = StartArgs( + data_type="bfloat16", + linear_att_ssm_data_type="float32", + mtp_step=2, + ssm_state_mode=mode, + ssm_replay_config={"cache_len": 8} if mode == "replay" else {}, + ) config = Qwen3NextResourceConfig.from_model_config(model_config, args) logical_layer_ids = tuple(range(len(layer_types))) (full_state,) = Qwen3NextTpPartModel.create_full_att_states( @@ -70,6 +78,10 @@ def test_hybrid_layer_mapping_and_mtp_states( model_logical_layer_ids=logical_layer_ids, device="meta", ) + assert bounded_state.ssm_slots_per_req == (3 if mode == "native" else 1) + assert (bounded_state.ssm_backend is None) == (mode == "native") + if mode == "replay": + assert bounded_state.ssm_backend.capacity == 8 manager = AttentionStateManager(request_capacity=1) manager.register_state(model_role="main", state=full_state) manager.register_state(model_role="main", state=bounded_state) diff --git a/unit_tests/common/att_resources/test_mtp_state_commit.py b/unit_tests/common/att_resources/test_mtp_state_commit.py index a50ff4fc71..c8fe301217 100644 --- a/unit_tests/common/att_resources/test_mtp_state_commit.py +++ b/unit_tests/common/att_resources/test_mtp_state_commit.py @@ -80,6 +80,7 @@ def test_commit_without_recurrent_state_and_invalid_role(): def test_verify_commits_main_states_for_fixed_and_dynamic_widths(widths: tuple[int, ...]): manager = make_manager("cuda") backend = SimpleNamespace( + max_draft_step=2, model=SimpleNamespace( att_state_manager=manager, req_manager=SimpleNamespace( @@ -87,7 +88,7 @@ def test_verify_commits_main_states_for_fixed_and_dynamic_widths(widths: tuple[i req_to_next_token_ids=torch.tensor([[-1, 10, 11]] * 3, device="cuda") ) ), - ) + ), ) reqs, starts, tokens, expected_mask = [], [], [], [] for req, width, accepted in zip((2, 0, 1), widths, (3, 1, 2)): @@ -107,3 +108,35 @@ def test_verify_commits_main_states_for_fixed_and_dynamic_widths(widths: tuple[i assert manager.main_kda_conv_fp32_ssm_fp32_state.mtp_state_index.tolist() == [0, 1, 2] assert manager.main_gdn_conv_fp32_ssm_fp32_state.mtp_state_index.tolist() == [0, 1, 2] assert manager.draft0_gdn_conv_fp32_ssm_fp32_state.mtp_state_index.tolist() == [0, 0, 0] + + +def test_verify_routes_linear_commit_context_and_other_bounded_states(monkeypatch): + from lightllm.server.router.model_infer.mtp_speculative import utils + + reqs = torch.tensor([2, 2, 0, 0], dtype=torch.int32) + starts = torch.tensor([0, 2], dtype=torch.int32) + positions = torch.tensor([0, 1, 0, 1], dtype=torch.int32) + lengths = torch.tensor([2, 1], dtype=torch.int32) + mask = torch.tensor([1, 1, 1, 0], dtype=torch.int32) + monkeypatch.setattr(utils, "mtp_verify", lambda **kwargs: (lengths, mask)) + calls = [] + linear = object.__new__(GDNState) + linear.commit_accepted_state = lambda **kwargs: calls.append(kwargs) + other = SimpleNamespace(commit_bounded_att_state=lambda indices, steps: calls.append((indices, steps))) + resources = SimpleNamespace(get_bounded_attention_states=lambda role: (linear, other) if role == "main" else ()) + backend = SimpleNamespace( + max_draft_step=1, + model=SimpleNamespace( + att_state_manager=resources, + req_manager=SimpleNamespace(req_sampling_params_manager=SimpleNamespace(req_to_next_token_ids=None)), + ), + ) + batches = ((2, object()),) + utils.verify_mtp_tokens(backend, torch.zeros(4), reqs, starts, positions, linear_state_batches=batches) + assert calls[0]["linear_state_batches"] is batches + assert calls[0]["b_req_idx"] is reqs + assert calls[0]["b_mtp_index"] is positions + assert calls[0]["accepted_index"] is mask + assert calls[0]["verify_width"] == 2 + assert calls[1][0].tolist() == [2, 0] + assert calls[1][1] is lengths diff --git a/unit_tests/common/basemodel/attention/linear/test_gdn.py b/unit_tests/common/basemodel/attention/linear/test_gdn.py index 8ec77c6f4d..a775cca4bc 100644 --- a/unit_tests/common/basemodel/attention/linear/test_gdn.py +++ b/unit_tests/common/basemodel/attention/linear/test_gdn.py @@ -38,7 +38,7 @@ def test_prefill_casts_final_state_to_cache_dtype(monkeypatch, cache_dtype): infer_state = InferStateInfo() infer_state.b1_cu_q_seq_len = torch.tensor([0, 1], dtype=torch.int32) infer_state.b_ready_cache_len = 0 - state = gdn.LinearAttPrefillAttState( + state = gdn.GDNLinearAttPrefillAttState( backend=backend, _infer_state_weak_ref=weakref.ref(infer_state), ) @@ -110,7 +110,7 @@ def test_decode_state_initializes_dynamic_mtp_layout(monkeypatch, role: str): def build_params(**kwargs): build_calls.append(kwargs) - return expected_cu_q_seq_len, expected_conv_buffer_idx, expected_num_accepted_tokens + return expected_cu_q_seq_len, expected_conv_buffer_idx monkeypatch.setattr(gdn, "build_dynamic_mtp_linear_att_state_params", build_params) state, infer_state = _create_decode_state( @@ -125,7 +125,7 @@ def build_params(**kwargs): assert state.b1_mtp_cu_q_seq_len is expected_cu_q_seq_len assert state.b_conv_buffer_idx is expected_conv_buffer_idx - assert state.b_num_accepted_tokens is expected_num_accepted_tokens + torch.testing.assert_close(state.b_num_accepted_tokens, expected_num_accepted_tokens) torch.testing.assert_close( state.b_ssm_buffer_idx, torch.tensor( @@ -136,12 +136,11 @@ def build_params(**kwargs): assert len(build_calls) == 1 assert build_calls[0]["b_req_idx"] is state.infer_state.b_req_idx assert build_calls[0]["b_mtp_index"] is state.infer_state.b_mtp_index - assert build_calls[0]["req_to_mtp_state_index"] is req_to_mtp_state_index assert build_calls[0]["hold_req_id"] == 5 def test_linear_backend_is_abstract_and_triton_supplies_chunk_kernel(): - assert inspect.isabstract(gdn.LinearAttBackend) + assert inspect.isabstract(gdn.GDNLinearAttBackend) backend = object.__new__(TritonLinearAttBackend) assert backend.get_prefill_kernel() is fla_ops.chunk_gated_delta_rule @@ -358,6 +357,8 @@ def _create_decode_state(*, draft_step, dynamic_layout, b_req_idx, req_to_mtp_st HOLD_REQUEST_ID=req_to_mtp_state_index.shape[0] - 1 if req_to_mtp_state_index is not None else -1, ) infer_state.att_state_role = role + if req_to_mtp_state_index is None: + req_to_mtp_state_index = torch.zeros(int(b_req_idx.max()) + 2, dtype=torch.int32) if req_to_mtp_state_index is not None: request_capacity = req_to_mtp_state_index.numel() resources = AttentionStateManager(request_capacity=request_capacity) @@ -379,7 +380,7 @@ def _create_decode_state(*, draft_step, dynamic_layout, b_req_idx, req_to_mtp_st resources.register_state(model_role, linear) resources.finish_state_registration() infer_state.att_state_manager = resources - state = gdn.LinearAttDecodeAttState( + state = gdn.GDNLinearAttDecodeAttState( backend=SimpleNamespace( model=model, uses_dynamic_spec_verify_layout=lambda: dynamic_layout, diff --git a/unit_tests/common/basemodel/test_cuda_graph_autotune.py b/unit_tests/common/basemodel/test_cuda_graph_autotune.py index b5b50beb20..71ff5368c8 100644 --- a/unit_tests/common/basemodel/test_cuda_graph_autotune.py +++ b/unit_tests/common/basemodel/test_cuda_graph_autotune.py @@ -142,11 +142,12 @@ def test_linear_config_is_local_to_capture_state(capture_env, overlap, kda): from lightllm.common.basemodel.hidden_collector import NoopHiddenCollector from lightllm.common.basemodel.infer_struct import InferStateInfo from lightllm.common.basemodel.triton_kernel.linear_att.compact_ssm import CompactSSMBackend - from lightllm.common.linear_state import LinearStateManager, SSMRunConfig + from lightllm.common.linear_state import SSMRunConfig + from lightllm.common.attention_state.bounded_att import GDNState, KDAState capture_env.args.enable_decode_microbatch_overlap = overlap graph = CudaGraph(1, 2, 1, max_batch_size=2) - manager = object.__new__(LinearStateManager) + manager = object.__new__(KDAState if kda else GDNState) manager.ssm_backend = object.__new__(CompactSSMBackend) att_class = KDALinearAttDecodeAttState if kda else GDNLinearAttDecodeAttState states = [] @@ -154,7 +155,7 @@ def test_linear_config_is_local_to_capture_state(capture_env, overlap, kda): for _ in range(2 if overlap else 1): state = InferStateInfo() state.input_ids = torch.ones(2, dtype=torch.int64) - state.req_manager = SimpleNamespace(linear_state_manager=manager) + state.att_state_manager = SimpleNamespace(get_bounded_attention_states=lambda role: (manager,)) state.hidden_collector = NoopHiddenCollector() state.decode_att_state = TritonDecodeAttState(_infer_state_weak_ref=weakref.ref(state)) state.decode_att_state1 = att_class(_infer_state_weak_ref=weakref.ref(state)) diff --git a/unit_tests/common/basemodel/triton_kernel/linear_att/test_acceptance_fusion.py b/unit_tests/common/basemodel/triton_kernel/linear_att/test_acceptance_fusion.py index 483f8d50d9..13ee79dcfb 100644 --- a/unit_tests/common/basemodel/triton_kernel/linear_att/test_acceptance_fusion.py +++ b/unit_tests/common/basemodel/triton_kernel/linear_att/test_acceptance_fusion.py @@ -5,7 +5,6 @@ from lightllm.common.basemodel.triton_kernel.mtp_utils import linear_att_mtp_state_index_update as update from lightllm.common.basemodel.triton_kernel.linear_att.replayssm import ReplaySSMBackend from lightllm.common.linear_state import LinearStateManager -from lightllm.common.req_manager.linear_att import ReqManagerForMamba def make(n, width): @@ -26,22 +25,19 @@ def make(n, width): ssm_backend.hold_req_id = hold ssm_backend.capacity = max(4, triton.next_power_of_2(width)) ssm_backend.verify_width = width - manager = ReqManagerForMamba.__new__(ReqManagerForMamba) - manager.linear_state_manager = object.__new__(LinearStateManager) - manager.linear_state_manager.accepted_state_index = tensor([-9] * (hold + 1)) - manager.linear_state_manager.ssm_backend = ssm_backend + manager = object.__new__(LinearStateManager) + manager.mtp_state_index = tensor([-9] * (hold + 1)) + manager.ssm_backend = ssm_backend return manager, args, (ids, lengths, starts) def accept_unfused(m, args): - update(m.linear_state_manager.accepted_state_index, *args) - m.linear_state_manager.ssm_backend.accept_updates( - args[1][args[0].long()], m.linear_state_manager.accepted_state_index - ) + update(m.mtp_state_index, *args) + m.ssm_backend.accept_updates(args[1][args[0].long()], m.mtp_state_index) def accept_fused(m, args): - m.update_mtp_state(*args) + m.commit_accepted_state(*args) @pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") @@ -58,19 +54,19 @@ def test_acceptance_fusion_matches_original_and_cpu_under_graph(): graph = torch.cuda.CUDAGraph() with torch.cuda.graph(graph): accept_fused(m, args) - expected_index = [-9] * len(m.linear_state_manager.accepted_state_index) + expected_index = [-9] * len(m.mtp_state_index) expected_cursor = [7] * len(expected_index) - m.linear_state_manager.accepted_state_index.fill_(-9) - ref.linear_state_manager.accepted_state_index.fill_(-9) - m.linear_state_manager.ssm_backend.cursors.fill_(7) - ref.linear_state_manager.ssm_backend.cursors.fill_(7) + m.mtp_state_index.fill_(-9) + ref.mtp_state_index.fill_(-9) + m.ssm_backend.cursors.fill_(7) + ref.ssm_backend.cursors.fill_(7) for step in range(24): mask = [] for i, (r, l) in enumerate(zip(ids, lengths)): count = (step + i) % (l + 1) mask.extend([1] * count + [0] * (l - count)) expected_index[r] = count - 1 - if r != m.linear_state_manager.ssm_backend.hold_req_id: + if r != m.ssm_backend.hold_req_id: expected_cursor[r] += count if width > 1 else 1 args[3].copy_(torch.tensor(mask, device="cuda", dtype=torch.int32)) accept_unfused(ref, args) @@ -78,19 +74,15 @@ def test_acceptance_fusion_matches_original_and_cpu_under_graph(): graph.replay() else: accept_fused(m, args) - assert torch.equal( - m.linear_state_manager.accepted_state_index, ref.linear_state_manager.accepted_state_index - ) - assert torch.equal( - m.linear_state_manager.ssm_backend.cursors, ref.linear_state_manager.ssm_backend.cursors - ) - assert m.linear_state_manager.accepted_state_index.tolist() == expected_index - assert m.linear_state_manager.ssm_backend.cursors.tolist() == expected_cursor + assert torch.equal(m.mtp_state_index, ref.mtp_state_index) + assert torch.equal(m.ssm_backend.cursors, ref.ssm_backend.cursors) + assert m.mtp_state_index.tolist() == expected_index + assert m.ssm_backend.cursors.tolist() == expected_cursor checks += 1 # Native optional-pointer path and manager dispatch leave metadata-only backend absent. - m.linear_state_manager.ssm_backend = None + m.ssm_backend = None accept_fused(m, args) - assert m.linear_state_manager.accepted_state_index.tolist() == expected_index + assert m.mtp_state_index.tolist() == expected_index # Non-Replay backends must still receive the accepted request IDs and indexes. class CompactProbe: @@ -99,10 +91,10 @@ def accept_updates(self, reqs, index): m, args, (ids, _, _) = make(7, 4) probe = CompactProbe() - m.linear_state_manager.ssm_backend = probe + m.ssm_backend = probe accept_fused(m, args) assert probe.seen[0].tolist() == ids - assert torch.equal(probe.seen[1], m.linear_state_manager.accepted_state_index) + assert torch.equal(probe.seen[1], m.mtp_state_index) print("ACCEPT_EQUIVALENCE_CHECKS=" + str(checks), flush=True) @@ -126,10 +118,9 @@ def test_fused_acceptance_preserves_folds_snapshots_and_slot_reuse(dtype, comput ] managers = [] for ssm_backend in ssm_backends: - manager = ReqManagerForMamba.__new__(ReqManagerForMamba) - manager.linear_state_manager = object.__new__(LinearStateManager) - manager.linear_state_manager.ssm_backend = ssm_backend - manager.linear_state_manager.accepted_state_index = torch.zeros(8, dtype=torch.int32, device="cuda") + manager = object.__new__(LinearStateManager) + manager.ssm_backend = ssm_backend + manager.mtp_state_index = torch.zeros(8, dtype=torch.int32, device="cuda") managers.append(manager) reqs = torch.full((4,), 7, dtype=torch.int32, device="cuda") cu = torch.tensor([0, 3, 6, 9, 12], dtype=torch.int32, device="cuda") @@ -188,9 +179,7 @@ def cycle(i): assert torch.equal(left, right) assert torch.equal(ssm_backends[0].state, ssm_backends[1].state) assert torch.equal(ssm_backends[0].cursors, ssm_backends[1].cursors) - assert torch.equal( - managers[0].linear_state_manager.accepted_state_index, managers[1].linear_state_manager.accepted_state_index - ) + assert torch.equal(managers[0].mtp_state_index, managers[1].mtp_state_index) if iteration % 7 == 0: assert torch.equal(ssm_backends[0].get_accepted_state(3), ssm_backends[1].get_accepted_state(3)) if iteration % 11 == 10: @@ -199,9 +188,9 @@ def cycle(i): assert torch.equal(ssm_backends[0].state, ssm_backends[1].state) if iteration == 19: for manager in managers: - manager.linear_state_manager.ssm_backend.clear_history(1) - manager.linear_state_manager.ssm_backend.state[:, 1].zero_() - manager.linear_state_manager.accepted_state_index[1] = 0 + manager.ssm_backend.clear_history(1) + manager.ssm_backend.state[:, 1].zero_() + manager.mtp_state_index[1] = 0 for ssm_backend in ssm_backends: ssm_backend.merge_accepted_updates(reqs[:3]) assert torch.equal(ssm_backends[0].state, ssm_backends[1].state) diff --git a/unit_tests/common/basemodel/triton_kernel/linear_att/test_mtp_state_params.py b/unit_tests/common/basemodel/triton_kernel/linear_att/test_mtp_state_params.py index 3c5ef837d3..32cba5c937 100644 --- a/unit_tests/common/basemodel/triton_kernel/linear_att/test_mtp_state_params.py +++ b/unit_tests/common/basemodel/triton_kernel/linear_att/test_mtp_state_params.py @@ -15,7 +15,7 @@ def test_dynamic_sequences_stay_within_request_capacity_under_graph(request_capa reqs = torch.full((token_bucket,), hold, dtype=torch.int32, device="cuda") mtp_index = torch.zeros_like(reqs) manager = object.__new__(LinearStateManager) - manager.accepted_state_index = torch.arange(request_capacity + 1, dtype=torch.int32, device="cuda") % 4 + manager.mtp_state_index = torch.arange(request_capacity + 1, dtype=torch.int32, device="cuda") % 4 manager.ssm_slots_per_req = 4 manager.ssm_backend = None sequences = min(request_capacity, token_bucket) 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 index ebbb915a15..165b86b131 100644 --- a/unit_tests/common/basemodel/triton_kernel/linear_att/test_replayssm.py +++ b/unit_tests/common/basemodel/triton_kernel/linear_att/test_replayssm.py @@ -549,130 +549,70 @@ def test_autotune_candidates_match_recurrence(dtype, run_config, mode): @pytest.mark.parametrize("save_big", [False, True]) @pytest.mark.parametrize("dtype", [torch.float32, torch.bfloat16]) -def test_checkpoint_restores_pending_history_and_accepted_conv(save_big, monkeypatch, dtype): - from types import SimpleNamespace - from lightllm.common.linear_state import LinearStateManager - 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 +def test_checkpoint_restores_pending_history_and_accepted_conv(save_big, dtype): + from lightllm.common.attention_state.bounded_att import ConvSSMStateConfig + from unit_tests.common.test_linear_state import make_state torch.manual_seed(42) state_rtol = 2e-5 if dtype == torch.float32 else torch.finfo(dtype).eps - # Real save/restore and PD helpers, with just the unrelated allocator omitted. - manager = object.__new__(ReqManagerForMamba) - manager.linear_state_manager = object.__new__(LinearStateManager) - config = SimpleNamespace( - linear_layer_num=2, - conv_state_dtype=torch.bfloat16, - ssm_state_dtype=dtype, - 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, + config = ConvSSMStateConfig( + logical_layer_ids=(0, 1), + global_k_heads=2, + global_v_heads=2, + key_dim=32, + value_dim=64, + conv_width=3, + conv_dtype=torch.bfloat16, + ssm_dtype=dtype, ) - manager.linear_config = config - manager.linear_state_manager.linear_config = config - manager.mtp_step = 2 - manager.linear_state_manager.ssm_slots_per_req = 1 - manager.linear_state_manager.accepted_state_index = torch.zeros(4, device="cuda", dtype=torch.int32) - manager.linear_state_manager.accepted_state_index[1] = 1 - manager.linear_state_manager.conv_state = LayerCache(4, torch.bfloat16, (192, 5), 2, "cuda") - manager.linear_state_manager.conv_state.buffer.copy_( - torch.randn_like(manager.linear_state_manager.conv_state.buffer) - ) - manager.linear_state_manager.ssm_state = LayerCache(4, dtype, (2, 32, 64), 2, "cuda") - manager.linear_state_manager.ssm_state.buffer.zero_() - manager.linear_state_manager.ssm_backend = ReplaySSMBackend(manager.linear_state_manager.ssm_state.buffer, 16, 3) - ssm_backend = manager.linear_state_manager.ssm_backend - ssm_backend.raw_keys.normal_() - ssm_backend.raw_values.normal_() - ssm_backend.gates.uniform_(-0.1, -0.01) - ssm_backend.betas.uniform_(0.1, 0.9) - ssm_backend.cursors[1] = 2 - active_before = manager.linear_state_manager.ssm_state.buffer[:, 1].clone() - expected = manager.linear_state_manager.ssm_state.buffer[:, 1].float().clone() - for i in range(2): - key = ssm_backend.raw_keys[:, 1, :, i].float() - key /= (key.square().sum(-1, keepdim=True) + 1e-6).sqrt() - expected *= ssm_backend.gates[:, 1, :, i].exp()[..., None, None] - delta = ssm_backend.betas[:, 1, :, i, None] * ( - ssm_backend.raw_values[:, 1, :, i].float() - torch.einsum("lhkv,lhk->lhv", expected, key) - ) - expected += key[..., None] * delta[..., None, :] - expected = expected.to(dtype) - conv_expected = manager.linear_state_manager.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) - materialized = ssm_backend.get_accepted_state(1) - torch.testing.assert_close(materialized, expected, rtol=state_rtol, atol=2e-5) - 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) - # Transfer/restore are exact; only the independent recurrence reference has - # a tolerance for FP32 reduction differences crossing BF16 rounding points. - torch.testing.assert_close(state_saved, materialized.cpu(), rtol=0, atol=0) - torch.testing.assert_close(conv_saved, conv_expected.cpu(), rtol=0, atol=0) - torch.testing.assert_close(manager.linear_state_manager.ssm_state.buffer[:, 1], active_before, rtol=0, atol=0) - assert ssm_backend.cursors[1].item() == 2 - ssm_backend.cursors[2] = 13 # stale history from a previous owner must disappear - manager.restore_state(SimpleNamespace(req_idx=2), cpu_cache, 0) - assert ssm_backend.cursors[2].item() == 0 - assert manager.linear_state_manager.accepted_state_index[2].item() == 0 - torch.testing.assert_close(manager.linear_state_manager.ssm_state.buffer[:, 2], materialized, rtol=0, atol=0) - torch.testing.assert_close( - manager.linear_state_manager.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, - linear_state_manager=manager.linear_state_manager, - ) - helper = Qwen3NextLinearAttPageHelper(mem) - # A second pending suffix exercises PD export, independently of CPU save. - ssm_backend.raw_keys[:, 1, :, 0].normal_() - ssm_backend.raw_values[:, 1, :, 0].normal_() - ssm_backend.gates[:, 1, :, 0].uniform_(-0.1, -0.01) - ssm_backend.betas[:, 1, :, 0].uniform_(0.1, 0.9) - ssm_backend.cursors[1] = 1 - key = ssm_backend.raw_keys[:, 1, :, 0].float() - key /= (key.square().sum(-1, keepdim=True) + 1e-6).sqrt() - expected = active_before.float().clone() - expected *= ssm_backend.gates[:, 1, :, 0].exp()[..., None, None] - delta = ssm_backend.betas[:, 1, :, 0, None] * ( - ssm_backend.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) - expected = expected.to(dtype) - ssm_page = torch.empty(helper.ssm_shape, device="cuda", dtype=dtype) - helper._write_one_rank(mem, 0, 1, conv_page, ssm_page) - torch.testing.assert_close(ssm_page, expected, rtol=state_rtol, atol=2e-5) - torch.testing.assert_close(manager.linear_state_manager.ssm_state.buffer[:, 1], active_before, rtol=0, atol=0) - assert ssm_backend.cursors[1].item() == 1 - ssm_backend.cursors[2] = 7 - manager.linear_state_manager.accepted_state_index[2] = 2 - helper._read_one_rank(mem, 0, 2, conv_page, ssm_page) - torch.testing.assert_close(manager.linear_state_manager.ssm_state.buffer[:, 2], ssm_page, rtol=0, atol=0) - torch.testing.assert_close( - manager.linear_state_manager.conv_state.buffer[:, 2, :, :3], conv_expected, rtol=0, atol=0 - ) - assert ssm_backend.cursors[2].item() == manager.linear_state_manager.accepted_state_index[2].item() == 0 - manager.init_hybrid_attention_state(SimpleNamespace(req_idx=2)) - assert torch.count_nonzero(manager.linear_state_manager.ssm_state.buffer[:, 2]).item() == 0 - assert torch.count_nonzero(manager.linear_state_manager.conv_state.buffer[:, 2]).item() == 0 + state, _, layout = make_state(mode="replay", config=config) + state.initialize_page_pools(big_page_capacity=2, small_page_capacity=2) + state.mtp_state_index[1] = 1 + state.conv_state.normal_() + state.ssm_state.zero_() + backend = state.ssm_backend + backend.raw_keys.normal_() + backend.raw_values.normal_() + backend.gates.uniform_(-0.1, -0.01) + backend.betas.uniform_(0.1, 0.9) + active_before = state.ssm_state[:, 1].clone() + conv_expected = state.conv_state[:, 1, :, 1:4].clone() + block = torch.empty((1, layout.byte_end), dtype=torch.uint8, device="cuda") + kind = "big" if save_big else "small" + pool = state.big_pages_pool if save_big else state.small_pages_pool + for count in (2, 1): + backend.cursors[1] = count + expected = active_before.float().clone() + for i in range(count): + key = backend.raw_keys[:, 1, :, i].float() + key /= (key.square().sum(-1, keepdim=True) + 1e-6).sqrt() + expected *= backend.gates[:, 1, :, i].exp()[..., None, None] + delta = backend.betas[:, 1, :, i, None] * ( + backend.raw_values[:, 1, :, i].float() - torch.einsum("lhkv,lhk->lhv", expected, key) + ) + expected += key[..., None] * delta[..., None, :] + materialized = backend.get_accepted_state(1) + torch.testing.assert_close(materialized, expected.to(dtype), rtol=state_rtol, atol=2e-5) + state.save_bounded_att_runtime_to_pages([1], [8], kind, [0]) + state.write_to_pd_page(1, block, 0, 8, layout) + torch.cuda.synchronize() + torch.testing.assert_close(pool.ssm_state[0], materialized.cpu(), rtol=0, atol=0) + torch.testing.assert_close(pool.conv_state[0], conv_expected.cpu(), rtol=0, atol=0) + torch.testing.assert_close(state.ssm_state[:, 1], active_before, rtol=0, atol=0) + assert backend.cursors[1].item() == count + for restore in ( + lambda: state.restore_bounded_att_runtime_from_pages([2], [8], kind, [0]), + lambda: state.read_from_pd_page(block, 2, 0, 8, layout), + ): + backend.cursors[2] = 3 + state.mtp_state_index[2] = 2 + restore() + assert backend.cursors[2].item() == state.mtp_state_index[2].item() == 0 + torch.testing.assert_close(state.ssm_state[:, 2], materialized, rtol=0, atol=0) + torch.testing.assert_close(state.conv_state[:, 2, :, :3], conv_expected, rtol=0, atol=0) + state.reset_bounded_att_state([2]) + assert not state.ssm_state[:, 2].any() + assert not state.conv_state[:, 2].any() @pytest.mark.parametrize("dtype", [torch.float32, torch.bfloat16]) diff --git a/unit_tests/common/basemodel/triton_kernel/linear_att/test_replayssm_execution.py b/unit_tests/common/basemodel/triton_kernel/linear_att/test_replayssm_execution.py index 0dac93dad9..e172b6495c 100644 --- a/unit_tests/common/basemodel/triton_kernel/linear_att/test_replayssm_execution.py +++ b/unit_tests/common/basemodel/triton_kernel/linear_att/test_replayssm_execution.py @@ -6,7 +6,6 @@ build_dynamic_mtp_linear_att_state_params, ) from lightllm.common.linear_state import LinearStateManager -from lightllm.common.req_manager.linear_att import ReqManagerForMamba pytestmark = pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") @@ -32,10 +31,9 @@ def test_request_manager_accept_with_replay_graph(dtype, capture_metadata, varle ) for cls in [ReplaySSMBackend, ReplaySSMBackend] ] - manager = object.__new__(ReqManagerForMamba) - manager.linear_state_manager = object.__new__(LinearStateManager) - manager.linear_state_manager.ssm_backend = ssm_backends[1] - manager.linear_state_manager.accepted_state_index = torch.zeros(5, device="cuda", dtype=torch.int32) + manager = object.__new__(LinearStateManager) + manager.ssm_backend = ssm_backends[1] + manager.mtp_state_index = torch.zeros(5, device="cuda", dtype=torch.int32) rows = torch.full((12,), 4, device="cuda", dtype=torch.int32) mtp = torch.zeros_like(rows) flags = torch.zeros_like(rows) @@ -82,7 +80,7 @@ def forward_accept(index, cu, reqs): if index == 0: ssm_backend.accept_updates(reqs, accepted) else: - manager.update_mtp_state(starts, rows, mtp, flags, 4) + manager.commit_accepted_state(starts, rows, mtp, flags, 4) return out inputs(0) @@ -92,7 +90,7 @@ def forward_accept(index, cu, reqs): forward_accept(index, cu, reqs) ssm_backend.state.copy_(initial) ssm_backend.cursors.zero_() - manager.linear_state_manager.accepted_state_index.zero_() + manager.mtp_state_index.zero_() g = torch.cuda.CUDAGraph() with torch.cuda.graph(g): if capture_metadata: @@ -104,7 +102,7 @@ def forward_accept(index, cu, reqs): for ssm_backend in ssm_backends: ssm_backend.state.copy_(initial) ssm_backend.cursors.zero_() - manager.linear_state_manager.accepted_state_index.zero_() + manager.mtp_state_index.zero_() folds = 0 for iteration in range(32): ids, lengths, counts, logical = inputs(iteration) @@ -128,7 +126,7 @@ def forward_accept(index, cu, reqs): torch.testing.assert_close( ssm_backends[1].state.float(), ssm_backends[0].state.float(), atol=0.002, rtol=torch.finfo(dtype).eps ) - table = manager.linear_state_manager.accepted_state_index.cpu().tolist() + table = manager.mtp_state_index.cpu().tolist() for req, count in zip(ids, counts): assert table[req] == count - 1 assert table[2] == table[4] == 0 diff --git a/unit_tests/common/basemodel/triton_kernel/linear_att/test_ssm_autotune.py b/unit_tests/common/basemodel/triton_kernel/linear_att/test_ssm_autotune.py index 123c2cc910..03c13144f6 100644 --- a/unit_tests/common/basemodel/triton_kernel/linear_att/test_ssm_autotune.py +++ b/unit_tests/common/basemodel/triton_kernel/linear_att/test_ssm_autotune.py @@ -325,17 +325,17 @@ def test_two_verify_graphs_accept_with_their_own_configs(monkeypatch, mode, proj from lightllm.common.basemodel.attention.linear.kda import KDALinearAttDecodeAttState from lightllm.common.basemodel.basemodel import TpPartBaseModel from lightllm.common.basemodel.batch_objs import PostLayerOutput, ModelMtpOutputCollector - from lightllm.common.req_manager.linear_att import ReqManagerForMamba inputs = make_inputs(mode, dtype, "cuda", compute_mode=projection, kda=kda) ssm_backend = inputs["ssm_backend"] initial = ssm_backend.state.clone() configs = [{"BV": 8, "num_warps": 1, "num_stages": 1}, {"BV": 32, "num_warps": 2, "num_stages": 1}] monkeypatch.setattr(tuning, "select_config", lambda *args: configs[args[2].shape[1] > 8]) - manager = object.__new__(ReqManagerForMamba) - manager.linear_state_manager = object.__new__(LinearStateManager) - manager.linear_state_manager.ssm_backend = ssm_backend - manager.linear_state_manager.accepted_state_index = torch.zeros(9, dtype=torch.int32, device="cuda") + from lightllm.common.attention_state.bounded_att import KDAState, GDNState + + manager = object.__new__(KDAState if kda else GDNState) + manager.ssm_backend = ssm_backend + manager.mtp_state_index = torch.zeros(9, dtype=torch.int32, device="cuda") buckets = [] for index, ids in enumerate(([0], [1, 3])): # The physical verify buckets include HOLD, but acceptance is unpadded. @@ -371,7 +371,7 @@ def forward(target): with torch.cuda.graph(graph): outputs, config = forward(ssm_backend) state = InferStateInfo() - state.req_manager = manager + state.att_state_manager = SimpleNamespace(get_bounded_attention_states=lambda role: (manager,)) state.hidden_collector = SimpleNamespace(finish_output=lambda **kw: ModelMtpOutputCollector()) att_class = KDALinearAttDecodeAttState if kda else GDNLinearAttDecodeAttState state.decode_att_state1 = att_class(_infer_state_weak_ref=weakref.ref(state), run_config=config) @@ -405,7 +405,7 @@ def forward(target): torch.testing.assert_close(outputs[layer], expected, atol=0, rtol=0) counts = [1 + (iteration + i) % 4 for i in range(3)] flags = torch.tensor([int(j < n) for n in counts for j in range(4)], dtype=torch.int32, device="cuda") - manager.update_mtp_state( + manager.commit_accepted_state( starts, rows, mtp, @@ -414,7 +414,7 @@ def forward(target): linear_state_batches=tuple((len(b[-1]), b[2].linear_state_commit) for b in buckets), ) for _, _, _, ref, _, reqs, _, ids in buckets: - ref.accept_updates(reqs[:-1], manager.linear_state_manager.accepted_state_index) + ref.accept_updates(reqs[:-1], manager.mtp_state_index) torch.testing.assert_close(ssm_backend.state[:, ids], ref.state[:, ids], atol=0, rtol=0) if mode == "replay": assert torch.equal(ssm_backend.cursors[ids], ref.cursors[ids]) diff --git a/unit_tests/common/test_hybrid_req_manager_pd_sharing.py b/unit_tests/common/test_hybrid_req_manager_pd_sharing.py new file mode 100644 index 0000000000..f634a833e7 --- /dev/null +++ b/unit_tests/common/test_hybrid_req_manager_pd_sharing.py @@ -0,0 +1,132 @@ +import multiprocessing +import traceback +from multiprocessing.reduction import ForkingPickler + +import pytest +import torch + +from lightllm.common.attention_state.bounded_att import GDNState, KDAState +from lightllm.server.router.model_infer.mode_backend.pd.p2p_fix import reduce_tensor +from unit_tests.common.test_linear_state import make_state + + +class _SerializationFailure: + def __reduce_ex__(self, protocol): + raise RuntimeError("injected serialization failure") + + +def _assert_received_manager_roundtrip(payload, state_name, mode, expected): + resources = ForkingPickler.loads(payload) + state = getattr(resources, state_name) + assert state.big_pages_pool is state.small_pages_pool is None + assert resources.full_att_allocator is None + if state.ssm_backend is not None: + assert state.ssm_backend.state is state.ssm_state + layout = resources.get_pd_operator().get_pd_bounded_att_page_layouts()[state_name] + block = torch.empty((1, layout.byte_end), dtype=torch.uint8, device="cuda") + state.write_to_pd_page(1, block, 0, 8, layout) + state.read_from_pd_page(block, 2, 0, 8, layout) + actual = state.get_accepted_state(2) + for tensor, values in zip(actual, expected): + torch.testing.assert_close(tensor, torch.tensor(values, dtype=tensor.dtype, device="cuda"), rtol=0, atol=0) + assert state.mtp_state_index[2] == 0 + assert state.mtp_state_index[1] == 1 + if mode == "replay": + assert state.ssm_backend.cursors[2] == 0 + assert state.ssm_backend.cursors[1] == 1 + torch.cuda.synchronize() + + +def _receive_pd_managers(connection): + torch.cuda.set_device(0) + while True: + message = connection.recv() + if message is None: + return + try: + _assert_received_manager_roundtrip(*message) + connection.send(None) + except Exception: + connection.send(traceback.format_exc()) + + +@pytest.fixture(scope="module") +def pd_receiver(): + context = multiprocessing.get_context("spawn") + sender, receiver = context.Pipe() + process = context.Process(target=_receive_pd_managers, args=(receiver,)) + process.start() + receiver.close() + try: + yield sender, process + finally: + if process.is_alive(): + sender.send(None) + process.join(timeout=30) + if process.is_alive(): + process.terminate() + process.join(timeout=5) + sender.close() + assert process.exitcode == 0 + + +@pytest.mark.parametrize("cls", [GDNState, KDAState]) +@pytest.mark.parametrize("mode", ["native", "compact", "replay"]) +def test_pd_export_copy_keeps_worker_pools_and_runtime_identity(cls, mode): + state, resources, _ = make_state(cls=cls, mode=mode, device="meta") + state.big_pages_pool, state.small_pages_pool = object(), object() + before = vars(state).copy() + exported = resources.get_pd_operator().create_ipc_copy() + copy = getattr(exported, state.state_name) + assert copy.big_pages_pool is copy.small_pages_pool is None + assert copy.conv_state is state.conv_state + assert copy.ssm_state is state.ssm_state + assert copy.ssm_backend is state.ssm_backend + assert copy.mtp_state_index is state.mtp_state_index + assert all(vars(state)[name] is value for name, value in before.items()) + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA IPC and pinned memory required") +@pytest.mark.parametrize("cls", [GDNState, KDAState]) +@pytest.mark.parametrize("mode", ["native", "replay"]) +@pytest.mark.parametrize("fail_serialization", [False, True]) +def test_shared_attention_state_roundtrip_preserves_worker_resources( + monkeypatch, pd_receiver, cls, mode, fail_serialization +): + state, resources, _ = make_state(cls=cls, mode=mode) + state.initialize_page_pools(big_page_capacity=2, small_page_capacity=2) + state.conv_state.copy_(torch.arange(5, dtype=torch.bfloat16, device="cuda")) + for slot in range(state.ssm_state.shape[1]): + state.ssm_state[:, slot].fill_(slot + 1) + state.mtp_state_index[1], state.mtp_state_index[2] = 1, 2 + if mode == "replay": + state.ssm_backend.raw_keys.zero_() + state.ssm_backend.raw_values.zero_() + state.ssm_backend.gates.fill_(-0.5) + state.ssm_backend.betas.zero_() + state.ssm_backend.cursors[1], state.ssm_backend.cursors[2] = 1, 3 + expected = [tensor.cpu().tolist() for tensor in state.get_accepted_state(1)] + before = vars(state).copy() + source_conv, source_ssm = state.conv_state[:, 1].clone(), state.ssm_state[:, state.ssm_slots_per_req].clone() + snapshots = [] + for pool in (state.big_pages_pool, state.small_pages_pool): + for tensor in (pool.conv_state, pool.ssm_state): + tensor.fill_(3) + snapshots.append((tensor, tensor.data_ptr(), tensor.clone())) + monkeypatch.setitem(ForkingPickler._extra_reducers, torch.Tensor, reduce_tensor) + exported = resources.get_pd_operator().create_ipc_copy() + if fail_serialization: + with pytest.raises(RuntimeError, match="injected serialization failure"): + ForkingPickler.dumps((_SerializationFailure(), exported)) + else: + sender, receiver = pd_receiver + sender.send((ForkingPickler.dumps(exported).tobytes(), state.state_name, mode, expected)) + assert sender.poll(30), f"PD receiver did not respond; exit code: {receiver.exitcode}" + error = sender.recv() + assert error is None, error + assert all(vars(state)[name] is value for name, value in before.items()) + torch.testing.assert_close(state.conv_state[:, 1], source_conv, rtol=0, atol=0) + torch.testing.assert_close(state.ssm_state[:, state.ssm_slots_per_req], source_ssm, rtol=0, atol=0) + for tensor, pointer, contents in snapshots: + assert tensor.data_ptr() == pointer + torch.testing.assert_close(tensor, contents, rtol=0, atol=0) diff --git a/unit_tests/common/test_linear_state.py b/unit_tests/common/test_linear_state.py index 3c0aa0c5df..55de82d4e1 100644 --- a/unit_tests/common/test_linear_state.py +++ b/unit_tests/common/test_linear_state.py @@ -1,114 +1,149 @@ -from types import SimpleNamespace +from dataclasses import replace +from importlib import import_module import pytest import torch -from lightllm.common.kv_cache_mem_manager.qwen3next_mem_manager import Qwen3NextLinearAttPageHelper -from lightllm.common.linear_state import LinearStateManager -from lightllm.common.req_manager.linear_att import ReqManagerForMamba -from lightllm.common.state_cache_manager import LinearAttCacheConfig, LinearAttCacheManager +from lightllm.common.attention_state import AttentionStateManager +from lightllm.common.attention_state.bounded_att import ConvSSMStateConfig, GDNState, KDAState -pytestmark = pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") +def make_state(cls=GDNState, mode="native", step=2, dtype=torch.float32, device="cuda", config=None): + config = config or ConvSSMStateConfig( + logical_layer_ids=(0, 1), + global_k_heads=1, + global_v_heads=1 if cls is KDAState else 2, + key_dim=32, + value_dim=32, + conv_width=3, + conv_dtype=torch.bfloat16, + ssm_dtype=dtype, + ) + state = cls(config, 4, mtp_step=step, device=device, ssm_state_mode=mode, ssm_replay_config={"cache_len": 4}) + resources = AttentionStateManager(request_capacity=4) + resources.register_state("main", state) + resources.finish_state_registration() + layout = resources.get_pd_operator().get_pd_bounded_att_page_layouts()[state.state_name] + return state, resources, layout +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") +@pytest.mark.parametrize("cls", [GDNState, KDAState]) @pytest.mark.parametrize("mode,step", [("native", 0), ("native", 2), ("compact", 2), ("replay", 0), ("replay", 2)]) @pytest.mark.parametrize("dtype", [torch.float32, torch.bfloat16]) -def test_checkpoint_and_pd_export_same_accepted_prefix_and_restore_reused_slot(mode, step, dtype): - config = LinearAttCacheConfig( - tp_world_size=1, - full_att_all_num_kv_heads=1, - full_att_dtype=torch.bfloat16, - full_att_num_kv_heads=1, - full_att_head_dim=32, - global_linear_k_heads=1, - global_linear_v_heads=2, - num_linear_k_heads=1, - num_linear_v_heads=2, - head_linear_k_dim=32, - head_linear_v_dim=32, - conv_kernel_size=4, - linear_layer_num=2, - conv_state_dtype=torch.bfloat16, - ssm_state_dtype=dtype, - full_attention_interval=3, - all_layer_num=3, - ) - args = SimpleNamespace(mtp_step=step, ssm_state_mode=mode, ssm_replay_config={"cache_len": 4}) - state = LinearStateManager(3, config, args) - manager = object.__new__(ReqManagerForMamba) - manager.linear_state_manager = state - cache = LinearAttCacheManager(3, config) - manager.mem_manager = SimpleNamespace(big_page_buffers=cache) - mem = SimpleNamespace(linear_config=config, linear_state_manager=state) - helper = Qwen3NextLinearAttPageHelper(mem) +def test_checkpoint_and_pd_export_same_accepted_prefix_and_restore_reused_slot(cls, mode, step, dtype): + state, _, layout = make_state(cls, mode, step, dtype) + state.initialize_page_pools(big_page_capacity=2, small_page_capacity=2) src, dst, offset = 1, 2, min(1, step) - width = config.get_mtp_conv_state_shape(step)[-1] - state.conv_state.buffer.copy_(torch.arange(width, device="cuda").to(torch.bfloat16)) - for slot in range(state.ssm_state.buffer.shape[1]): - state.ssm_state.buffer[:, slot].fill_(slot + 1) - if state.accepted_state_index is not None: - state.accepted_state_index[src] = offset - expected_conv = state.conv_state.buffer[:, src, ..., offset : offset + 3].clone() - expected_ssm = state.ssm_state.buffer[ - :, src * state.ssm_slots_per_req + (offset if mode == "native" else 0) - ].clone() + state.conv_state.copy_(torch.arange(3 + step, device="cuda").to(torch.bfloat16)) + for slot in range(state.ssm_state.shape[1]): + state.ssm_state[:, slot].fill_(slot + 1) + if state.mtp_state_index is not None: + state.mtp_state_index[src] = offset + expected_conv = state.conv_state[:, src, ..., offset : offset + 3].clone() + expected_ssm = state.ssm_state[:, src * state.ssm_slots_per_req + (offset if mode == "native" else 0)].clone() if mode == "replay": - # One pending accepted update: beta=0 leaves only decay of the base. state.ssm_backend.raw_keys.zero_() state.ssm_backend.raw_values.zero_() state.ssm_backend.gates.fill_(-0.5) state.ssm_backend.betas.zero_() state.ssm_backend.cursors[src] = 1 expected_ssm = (expected_ssm.float() * torch.exp(torch.tensor(-0.5))).to(dtype) - before = state.ssm_state.buffer.clone() - snapshot = state.get_accepted_state(src) - torch.testing.assert_close(snapshot.ssm, expected_ssm, rtol=2e-5, atol=2e-5) - expected_ssm = snapshot.ssm.clone() - - manager.save_state(src, 0, cache) - manager.save_big_page_states(torch.tensor([src], dtype=torch.int32, device="cuda"), [src], [1]) - conv_page = torch.empty(helper.conv_shape, device="cuda", dtype=torch.bfloat16) - ssm_page = torch.empty(helper.ssm_shape, device="cuda", dtype=dtype) - helper._write_one_rank(mem, 0, src, conv_page, ssm_page) + before = state.ssm_state.clone() + block = torch.empty((1, layout.byte_end), dtype=torch.uint8, device="cuda") + state.write_to_pd_page(src, block, 0, 8, layout) + for kind in ("small", "big"): + state.save_bounded_att_runtime_to_pages([src], [8], kind, [0]) torch.cuda.synchronize() - for index in (0, 1): - conv, ssm = cache.get_state_cache(index) - torch.testing.assert_close(conv, expected_conv.cpu(), rtol=0, atol=0) - torch.testing.assert_close(ssm, expected_ssm.cpu(), rtol=0, atol=0) - torch.testing.assert_close(conv_page, expected_conv, rtol=0, atol=0) - torch.testing.assert_close(ssm_page, expected_ssm, rtol=0, atol=0) - assert torch.equal(state.ssm_state.buffer, before) + for pool in (state.small_pages_pool, state.big_pages_pool): + torch.testing.assert_close(pool.conv_state[0], expected_conv.cpu(), rtol=0, atol=0) + torch.testing.assert_close(pool.ssm_state[0], expected_ssm.cpu(), rtol=2e-5, atol=2e-5) + assert torch.equal(state.ssm_state, before) if mode == "replay": assert state.ssm_backend.cursors[src].item() == 1 - def dirty_destination(): - if state.accepted_state_index is not None: - state.accepted_state_index[dst] = step - if mode == "replay": - state.ssm_backend.cursors[dst] = 3 - for restore in ( - lambda: manager.restore_state(SimpleNamespace(req_idx=dst), cache, 0), - lambda: helper._read_one_rank(mem, 0, dst, conv_page, ssm_page), + lambda: state.restore_bounded_att_runtime_from_pages([dst], [8], "small", [0]), + lambda: state.restore_bounded_att_runtime_from_pages([dst], [8], "big", [0]), + lambda: state.read_from_pd_page(block, dst, 0, 8, layout), ): - dirty_destination() + if state.mtp_state_index is not None: + state.mtp_state_index[dst] = step + if mode == "replay": + state.ssm_backend.cursors[dst] = 3 restore() actual = state.get_accepted_state(dst) torch.testing.assert_close(actual.conv, expected_conv, rtol=0, atol=0) - torch.testing.assert_close(actual.ssm, expected_ssm, rtol=0, atol=0) - if state.accepted_state_index is not None: - assert state.accepted_state_index[dst].item() == 0 + torch.testing.assert_close(actual.ssm, expected_ssm, rtol=2e-5, atol=2e-5) + if state.mtp_state_index is not None: + assert state.mtp_state_index[dst].item() == 0 if mode == "replay": assert state.ssm_backend.cursors[dst].item() == 0 - - # A skipped checkpoint must not overwrite another cache owner. - untouched = tuple(x.clone() for x in cache.get_state_cache(2)) - manager.save_big_page_states(torch.tensor([src], dtype=torch.int32, device="cuda"), [src], [-1]) - torch.cuda.synchronize() - for actual, expected in zip(cache.get_state_cache(2), untouched): - assert torch.equal(actual, expected) - state.reset(dst) - assert not state.conv_state.buffer[:, dst].any() + state.reset_bounded_att_state([dst]) + assert not state.conv_state[:, dst].any() start = dst * state.ssm_slots_per_req - assert not state.ssm_state.buffer[:, start : start + state.ssm_slots_per_req].any() + assert not state.ssm_state[:, start : start + state.ssm_slots_per_req].any() + + +def copy_on_cpu(*, copies): + for source, destination in copies: + destination.copy_(source) + + +def test_pd_block_roundtrip_preserves_replicated_heads_and_accepted_mtp_state(monkeypatch): + # Exercise layout and accepted-state selection independently of the CUDA copy kernel. + monkeypatch.setattr(import_module(GDNState.__module__), "copy_bounded_att_tensors", copy_on_cpu) + config = ConvSSMStateConfig( + logical_layer_ids=(0,), + global_k_heads=1, + global_v_heads=2, + key_dim=2, + value_dim=3, + conv_width=3, + conv_dtype=torch.float32, + ssm_dtype=torch.float32, + tp_world_size=4, + ) + states, layouts, snapshots = [], [], [] + src, dst = 1, 2 + for rank in range(4): + state, _, layout = make_state(device="cpu", config=replace(config, tp_rank=rank)) + state.conv_state.fill_(-1) + state.ssm_state.fill_(-2) + state.mtp_state_index.fill_(1) + state.conv_state[:, src, :2, 1:4].fill_(10 + rank) + state.conv_state[:, src, 2:4, 1:4].fill_(20 + rank) + state.conv_state[:, src, 4:, 1:4].fill_(30 + rank) + state.ssm_state[:, src * state.ssm_slots_per_req + 1].fill_(40 + rank) + states.append(state) + layouts.append(layout) + snapshots.append((state.conv_state.clone(), state.ssm_state.clone())) + size = layouts[0].byte_end + block = torch.full((1, size + 16), 0xA5, dtype=torch.uint8) + for state, layout in zip(states, layouts): + state.write_to_pd_page(src, block, 0, 8, layout) + expected_conv = torch.cat( + [torch.full((1, channels, 3), value) for channels, value in ((2, 10), (2, 20), (3, 30), (3, 32))], + dim=1, + ).float() + expected_ssm = torch.stack((torch.full((1, 2, 3), 40), torch.full((1, 2, 3), 42)), dim=1).float() + conv_bytes = expected_conv.numel() * 4 + torch.testing.assert_close(block[0, :conv_bytes].view(torch.float32).view_as(expected_conv), expected_conv) + torch.testing.assert_close(block[0, conv_bytes:size].view(torch.float32).view_as(expected_ssm), expected_ssm) + assert torch.all(block[0, size:] == 0xA5) + exported = block.clone() + other_requests = torch.tensor([0, 1, 3]) + for rank, (state, layout, (conv_before, ssm_before)) in enumerate(zip(states, layouts, snapshots)): + torch.testing.assert_close(state.conv_state, conv_before) + torch.testing.assert_close(state.ssm_state, ssm_before) + state.read_from_pd_page(block, dst, 0, 8, layout) + expected_local_conv = torch.cat( + (expected_conv[:, :4], expected_conv[:, 4 + rank // 2 * 3 : 7 + rank // 2 * 3]), dim=1 + ) + torch.testing.assert_close(state.conv_state[:, dst, ..., :3], expected_local_conv) + torch.testing.assert_close(state.ssm_state[:, dst * 3], expected_ssm[:, rank // 2 : rank // 2 + 1]) + assert state.mtp_state_index.tolist() == [1, 1, 0, 1] + torch.testing.assert_close(state.conv_state[:, other_requests], conv_before[:, other_requests]) + other_slots = (other_requests[:, None] * 3 + torch.arange(3)).flatten() + torch.testing.assert_close(state.ssm_state[:, other_slots], ssm_before[:, other_slots]) + assert torch.equal(block, exported) diff --git a/unit_tests/common/test_linear_state_contract.py b/unit_tests/common/test_linear_state_contract.py index 32b1f98c33..217f27a0cb 100644 --- a/unit_tests/common/test_linear_state_contract.py +++ b/unit_tests/common/test_linear_state_contract.py @@ -6,7 +6,8 @@ import torch from lightllm.common.linear_state import LinearStateManager, SSMRunConfig -from lightllm.common.state_cache_manager import Glm5NextCacheConfig, LinearAttCacheConfig +from lightllm.common.attention_state.bounded_att import ConvSSMStateConfig, GDNState +from lightllm.models.glm5_next.config import Glm5NextResourceConfig from lightllm.server.core.objs.start_args_type import StartArgs @@ -30,49 +31,38 @@ def test_config_rejects_unrepresented_kernel_options(): @pytest.fixture def linear_config(): - return LinearAttCacheConfig( - tp_world_size=1, - full_att_all_num_kv_heads=1, - full_att_dtype=torch.bfloat16, - full_att_num_kv_heads=1, - full_att_head_dim=32, - global_linear_k_heads=1, - global_linear_v_heads=2, - num_linear_k_heads=1, - num_linear_v_heads=2, - head_linear_k_dim=32, - head_linear_v_dim=32, - conv_kernel_size=4, - linear_layer_num=6, - conv_state_dtype=torch.bfloat16, - ssm_state_dtype=torch.float32, - full_attention_interval=3, - all_layer_num=8, + return ConvSSMStateConfig( + logical_layer_ids=(0, 1, 3, 4, 6, 7), + global_k_heads=1, + global_v_heads=2, + key_dim=32, + value_dim=32, + conv_width=3, + conv_dtype=torch.bfloat16, + ssm_dtype=torch.float32, ) @pytest.mark.parametrize("layer,row", [(0, 0), (1, 1), (3, 2), (4, 3), (6, 4), (7, 5)]) def test_layer_buffers_follow_packed_linear_layers(linear_config, layer, row): - owner = object.__new__(LinearStateManager) - owner.linear_config = linear_config - conv = torch.arange(6).view(6, 1) - ssm = conv + 10 - owner.conv_state = SimpleNamespace(buffer=conv) - owner.ssm_state = SimpleNamespace(buffer=ssm) - assert linear_config.get_linear_layer_index(layer) == row - actual_conv, actual_ssm = owner.get_layer_buffers(layer) - assert actual_conv.item() == row - assert actual_ssm.item() == row + 10 + owner = GDNState(linear_config, request_capacity=1, device="cpu") + for i in range(6): + owner.conv_state[i].fill_(i) + owner.ssm_state[i].fill_(i + 10) + actual_conv, actual_ssm = owner.get_layer_view(layer) + assert (actual_conv == row).all() + assert (actual_ssm == row + 10).all() @pytest.mark.parametrize("layer", [-1, 2, 5, 8]) def test_linear_mapping_rejects_full_attention_and_out_of_range_layers(linear_config, layer): - with pytest.raises(AssertionError): - linear_config.get_linear_layer_index(layer) + owner = GDNState(linear_config, request_capacity=1, device="cpu") + with pytest.raises(KeyError): + owner.get_layer_view(layer) def test_model_configuration_retains_nondefault_kda_gate_bound(): - config = Glm5NextCacheConfig.from_model_config( + config = Glm5NextResourceConfig.from_model_config( { "kv_lora_rank": 512, "num_hidden_layers": 4, @@ -95,9 +85,7 @@ def test_model_configuration_retains_nondefault_kda_gate_bound(): @pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") @pytest.mark.parametrize("mode", ["compact", "replay"]) def test_kda_backend_receives_gate_bound_at_construction(linear_config, mode): - linear_config.gate_lower_bound = -2.5 - args = SimpleNamespace(mtp_step=2, ssm_state_mode=mode, ssm_replay_config={}) - owner = LinearStateManager(2, linear_config, args, kda=True) + owner = LinearStateManager(linear_config, 3, mtp_step=2, ssm_state_mode=mode, kda=True, gate_lower_bound=-2.5) assert owner.ssm_backend.gate_lower_bound == -2.5 @@ -126,7 +114,7 @@ def test_compact_commits_each_unpadded_partition_with_its_verify_config(monkeypa owner = object.__new__(LinearStateManager) owner.ssm_backend = object.__new__(CompactSSMBackend) - owner.accepted_state_index = torch.tensor([0, 1, 2, 0], dtype=torch.int32) + owner.mtp_state_index = torch.tensor([0, 1, 2, 0], dtype=torch.int32) configs = [SSMRunConfig(8, 1), SSMRunConfig(32, 2)] commits = [owner.create_commit(config) for config in configs] owner.ssm_backend.run_config = {"BV": 128, "num_warps": 4} @@ -146,5 +134,5 @@ def test_compact_commits_each_unpadded_partition_with_its_verify_config(monkeypa linear_state_batches=((1, commits[0]), (0, None), (2, commits[1])), ) assert [reqs for reqs, _, _ in calls] == [[2], [0, 1]] - assert all(accepted is owner.accepted_state_index for _, accepted, _ in calls) + assert all(accepted is owner.mtp_state_index for _, accepted, _ in calls) assert [config for _, _, config in calls] == configs diff --git a/unit_tests/common/test_paged_kv_transfer.py b/unit_tests/common/test_paged_kv_transfer.py deleted file mode 100644 index 8a057daec1..0000000000 --- a/unit_tests/common/test_paged_kv_transfer.py +++ /dev/null @@ -1,45 +0,0 @@ -from types import SimpleNamespace - -import pytest -import torch - -from lightllm.common.kv_cache_mem_manager.mem_manager import MemoryManager - - -pytestmark = pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA is required for page transfer") - - -@pytest.mark.parametrize("dtype", [torch.bfloat16, torch.uint8]) -@pytest.mark.parametrize("tp_size", [1, 2, 4]) -def test_read_page_reuses_index_buffer_on_copy_stream(dtype, tp_size): - page_size, layer_num, head_dim = 32, 3, 64 - manager = MemoryManager.__new__(MemoryManager) - manager.kv_move_buffer = torch.empty((2, page_size, layer_num, 2 * tp_size, head_dim), dtype=dtype, device="cuda") - manager._buffer_mem_indexes_tensors = [torch.empty(page_size, dtype=torch.int64, pin_memory=True) for _ in range(2)] - peers = [ - SimpleNamespace(kv_buffer=torch.zeros((layer_num, 96, 2, head_dim), dtype=dtype, device="cuda")) - for _ in range(tp_size) - ] - stream = torch.cuda.Stream() - for iteration, count in enumerate((5, 32, 1, 32, 5, 1)): - page_index = iteration % 2 - indexes = torch.randperm(96)[:count].tolist() - payload = torch.randint(0, 128, manager.kv_move_buffer[page_index].shape, device="cuda").to(dtype) - for peer in peers: - peer.kv_buffer.zero_() - stream.wait_stream(torch.cuda.current_stream()) - with torch.cuda.stream(stream): - manager.kv_move_buffer[page_index].copy_(payload) - manager.read_page_kv_move_buffer_to_mem( - indexes, page_index, dp_index=0, mem_managers=peers, dp_world_size=tp_size - ) - done = torch.cuda.Event() - done.record(stream) - # The service returns the page to its pool after this completion event. - done.synchronize() - selection = torch.tensor(indexes, dtype=torch.int64, device="cuda") - for rank, peer in enumerate(peers): - expected = torch.zeros_like(peer.kv_buffer) - values = torch.stack((payload[:count, :, rank], payload[:count, :, tp_size + rank]), dim=2) - expected.index_copy_(1, selection, values.permute(1, 0, 2, 3)) - assert torch.equal(peer.kv_buffer, expected) diff --git a/unit_tests/models/glm5_next/test_cache.py b/unit_tests/models/glm5_next/test_cache.py index 530b95436a..f07e853f20 100644 --- a/unit_tests/models/glm5_next/test_cache.py +++ b/unit_tests/models/glm5_next/test_cache.py @@ -48,12 +48,16 @@ def test_nope_cache_config_uses_native_mla_width(): @pytest.mark.parametrize("small_page", [False, True]) @pytest.mark.parametrize("tp_world_size", [1, 4]) @pytest.mark.parametrize("mtp_step", [0, 2]) +@pytest.mark.parametrize("mode", ["native", "compact", "replay"]) def test_hybrid_checkpoint_restore_and_packed_kv_copy( monkeypatch, small_page, tp_world_size, mtp_step, + mode, ): + if mode == "compact" and mtp_step == 0: + pytest.skip("compact requires MTP") monkeypatch.setenv("LIGHTLLM_CURRENT_RANK_IN_NODE", "0") monkeypatch.setenv("LIGHTLLM_CURRENT_DEVICE_ID", "0") monkeypatch.setattr("lightllm.common.req_manager.req_sampling_params.get_vocab_size", lambda _: 128) @@ -106,6 +110,7 @@ def test_hybrid_checkpoint_restore_and_packed_kv_copy( config=config, request_capacity=4, mtp_step=mtp_step, + ssm_state_mode=mode, full_state=main_full, full_allocator=FullAttKVAllocator(full.capacity, full.page_size), draft_full_state=draft_full, @@ -143,6 +148,16 @@ def test_hybrid_checkpoint_restore_and_packed_kv_copy( linear.ssm_state[:, 0].normal_() conv = linear.conv_state[:, 0, ..., :3].clone() ssm = linear.ssm_state[:, 0].clone() + if mode == "replay": + linear.ssm_backend.raw_keys.zero_() + linear.ssm_backend.raw_values.zero_() + linear.ssm_backend.gates.fill_(-0.5) + linear.ssm_backend.betas.zero_() + linear.ssm_backend.cursors[0] = 1 + ssm = linear.ssm_backend.get_accepted_state(0).clone() + linear.ssm_backend.cursors[2] = 3 + if mtp_step: + linear.mtp_state_index[2] = mtp_step resources.save_bounded_att_runtime_to_pages( req_indices=[0], state_token_lens=[8], @@ -162,12 +177,16 @@ def test_hybrid_checkpoint_restore_and_packed_kv_copy( ) torch.cuda.synchronize() assert torch.equal(linear.conv_state[:, 2, ..., :3], conv) - assert torch.equal(linear.ssm_state[:, 2 * (mtp_step + 1)], ssm) + assert torch.equal(linear.ssm_state[:, 2 * linear.ssm_slots_per_req], ssm) + if mode == "replay": + assert linear.ssm_backend.cursors[2] == 0 + if mtp_step: + assert linear.mtp_state_index[2] == 0 # 整组边界不依赖旧窗口，恢复大小页无需改动无重叠窗口。 assert torch.equal(tail.buffer, unchanged) resources.reset_bounded_att_state([2]) assert not linear.conv_state[:, 2].any() - assert not linear.ssm_state[:, 2 * (mtp_step + 1) : 3 * (mtp_step + 1)].any() + assert not linear.ssm_state[:, 2 * linear.ssm_slots_per_req : 3 * linear.ssm_slots_per_req].any() assert not tail.buffer[:, 2, :, : tail.config.candidate_dim].any() assert torch.isneginf(tail.buffer[:, 2, :, tail.config.candidate_dim :]).all() assert torch.equal(tail.buffer[:, [0, 1, 3]], unchanged[:, [0, 1, 3]]) diff --git a/unit_tests/models/glm5_next/test_mtp.py b/unit_tests/models/glm5_next/test_mtp.py index 82e029c5dd..2aa33161f5 100644 --- a/unit_tests/models/glm5_next/test_mtp.py +++ b/unit_tests/models/glm5_next/test_mtp.py @@ -127,17 +127,18 @@ def prefill(values, destination, tail): @pytest.mark.parametrize("dynamic", [False, True]) @pytest.mark.parametrize("role", ["main", "draft0"]) -def test_kda_backend_routes_accepted_conv_and_ssm_states(dynamic, role: str): +@pytest.mark.parametrize("mode,mtp_step", [("native", 2), ("compact", 2), ("replay", 2), ("native", 0), ("replay", 0)]) +def test_kda_backend_routes_accepted_conv_and_ssm_states(dynamic, role: str, mode, mtp_step): from types import SimpleNamespace from lightllm.common.basemodel.attention.base_att import AttControl - from lightllm.common.basemodel.attention.linear.kda import KDADecodeAttState + from lightllm.common.basemodel.attention.linear.kda import KDALinearAttDecodeAttState from lightllm.common.basemodel.triton_kernel.linear_att.causal_conv1d import causal_conv1d_update torch.manual_seed(53) - heads, dim, width = 2, 128, 3 + heads, dim, width = 2, 128, mtp_step + 1 hidden = heads * dim - lengths = [2, 1, 3] if dynamic else [3, 3, 3] + lengths = [2, 1, 3] if dynamic and mtp_step else [width] * 3 requests = [2, 0, 1] request_rows = [req for req, length in zip(requests, lengths) for _ in range(length)] offsets = [offset for length in lengths for offset in range(length)] @@ -145,10 +146,10 @@ def test_kda_backend_routes_accepted_conv_and_ssm_states(dynamic, role: str): request_rows += [3, 3] offsets += [0, 0] tokens = len(request_rows) - conv = torch.randn(4, 3 * hidden, 5, device="cuda", dtype=torch.bfloat16) + conv = torch.randn(4, 3 * hidden, 3 + mtp_step, device="cuda", dtype=torch.bfloat16) ssm = torch.randn(4 * width, heads, dim, dim, device="cuda") * 0.01 old_conv, expected_ssm = conv.clone(), ssm.clone() - accepted_offsets = torch.tensor([1, 2, 0, 0], dtype=torch.int32, device="cuda") + accepted_offsets = torch.tensor([1, 2, 0, 0] if mtp_step else [0, 0, 0, 0], dtype=torch.int32, device="cuda") mixed = torch.randn(tokens, 3 * hidden, device="cuda", dtype=torch.bfloat16) gate = torch.randn(tokens, hidden, device="cuda", dtype=torch.bfloat16) beta = torch.randn(tokens, heads, device="cuda", dtype=torch.bfloat16) @@ -186,7 +187,7 @@ def test_kda_backend_routes_accepted_conv_and_ssm_states(dynamic, role: str): expected_ssm[req * width + token] = state[0] row += 1 backend = SimpleNamespace( - mtp_step=2, + mtp_step=mtp_step, tp_num_heads=heads, head_dim=dim, tp_hidden_size=hidden, @@ -205,11 +206,19 @@ def test_kda_backend_routes_accepted_conv_and_ssm_states(dynamic, role: str): conv_dtype=conv.dtype, ssm_dtype=ssm.dtype, ) + selected = None for model_role in ("main", "draft0"): - linear = KDAState(config, request_capacity=4, mtp_step=2, device="meta") - linear.mtp_state_index = accepted_offsets if model_role == role else torch.zeros_like(accepted_offsets) + linear = KDAState(config, request_capacity=4, mtp_step=mtp_step, device="cuda", ssm_state_mode=mode) + if mtp_step: + linear.mtp_state_index.copy_(accepted_offsets if model_role == role else torch.zeros_like(accepted_offsets)) if model_role == role: - linear.conv_state, linear.ssm_state = conv.unsqueeze(0), ssm.unsqueeze(0) + linear.conv_state[0].copy_(conv) + if mode == "native": + linear.ssm_state[0].copy_(ssm) + else: + indexes = torch.arange(4, device="cuda") * width + accepted_offsets + linear.ssm_state[0].copy_(ssm[indexes]) + selected = linear resources.register_state(model_role, linear) resources.finish_state_registration() manager = SimpleNamespace(HOLD_REQUEST_ID=3) @@ -225,7 +234,7 @@ def test_kda_backend_routes_accepted_conv_and_ssm_states(dynamic, role: str): linear_A_log=SimpleNamespace(weight=a), linear_dt_bias=SimpleNamespace(weight=bias), ) - state = KDADecodeAttState(backend=backend, _infer_state_weak_ref=weakref.ref(infer)) + state = KDALinearAttDecodeAttState(backend=backend, _infer_state_weak_ref=weakref.ref(infer)) state.init_state() actual = state.decode_att( None, @@ -239,6 +248,23 @@ def test_kda_backend_routes_accepted_conv_and_ssm_states(dynamic, role: str): ), ) torch.testing.assert_close(actual.reshape(tokens, heads, dim)[:row], torch.cat(expected), atol=2e-3, rtol=1e-2) - torch.testing.assert_close(ssm, expected_ssm, atol=2e-3, rtol=1e-2) + if mode == "native": + torch.testing.assert_close(selected.ssm_state[0], expected_ssm, atol=2e-3, rtol=1e-2) + else: + ids = torch.tensor(requests, dtype=torch.int32, device="cuda") + if mtp_step: + accepted = torch.zeros(4, dtype=torch.int32, device="cuda") + accepted[ids] = torch.tensor(lengths, dtype=torch.int32, device="cuda") - 1 + if mode == "compact": + selected.ssm_backend.accept_updates(ids, accepted, run_config=state.run_config) + else: + selected.ssm_backend.accept_updates(ids, accepted) + for req, length in zip(requests, lengths): + torch.testing.assert_close( + selected.ssm_backend.get_accepted_state(req)[0], + expected_ssm[req * width + length - 1], + atol=2e-3, + rtol=1e-2, + ) if dynamic: - assert torch.equal(conv[3], old_conv[3]) + assert torch.equal(selected.conv_state[0, 3], old_conv[3]) diff --git a/unit_tests/models/glm5_next/test_pd_cache.py b/unit_tests/models/glm5_next/test_pd_cache.py index ff93312516..0a9c871a75 100644 --- a/unit_tests/models/glm5_next/test_pd_cache.py +++ b/unit_tests/models/glm5_next/test_pd_cache.py @@ -46,7 +46,7 @@ def make_mems(monkeypatch, request): get_env_start_args.cache_clear() set_env_start_args(dataclasses.asdict(args)) - def create(tp, target_mtp_step=None): + def create(tp, target_mtp_step=None, mode="native"): local_mtp_step = mtp_step if target_mtp_step is None else target_mtp_step config = Glm5NextResourceConfig( tp_world_size=tp, @@ -87,6 +87,7 @@ def create(tp, target_mtp_step=None): config=config, request_capacity=4, mtp_step=local_mtp_step, + ssm_state_mode=mode, tp_rank=rank, full_state=main, full_allocator=FullAttKVAllocator(full.capacity, full.page_size), @@ -112,16 +113,23 @@ def create(tp, target_mtp_step=None): @pytest.mark.parametrize("prefill_tp,decode_tp", [(1, 1), (4, 4), (1, 4), (4, 1)]) @pytest.mark.parametrize("remainder", range(4)) @pytest.mark.parametrize("decode_mtp_step", [0, 1, 5]) +@pytest.mark.parametrize("source_mode", ["native", "replay"]) +@pytest.mark.parametrize("dest_mode", ["native", "replay", "compact"]) def test_pd_kv_and_runtime_state_roundtrip( make_mems, prefill_tp, decode_tp, remainder, decode_mtp_step, + source_mode, + dest_mode, ): - source, dest = make_mems(prefill_tp), make_mems(decode_tp, target_mtp_step=decode_mtp_step) + if dest_mode == "compact" and decode_mtp_step == 0: + pytest.skip("compact requires MTP") + source = make_mems(prefill_tp, mode=source_mode) + dest = make_mems(decode_tp, target_mtp_step=decode_mtp_step, mode=dest_mode) + dst_ssm_stride = dest[0].linear.ssm_slots_per_req mtp_size = get_env_start_args().mtp_step + 1 - dst_mtp_size = decode_mtp_step + 1 layers = sum(len(state.config.logical_layer_groups) for state in source[0].full_states) src_req, dst_req = 0, 2 length = 8 + remainder @@ -159,7 +167,15 @@ def test_pd_kv_and_runtime_state_roundtrip( state.storage.view(torch.uint8)[:, src_indexes] = packed[offset:end] offset = end mem.linear.conv_state[:, src_req, ..., :3].copy_(conv[:, :, heads].reshape(6, -1, 3)) - mem.linear.ssm_state[:, src_req * mtp_size].copy_(ssm[:, heads]) + mem.linear.ssm_state[:, src_req * mem.linear.ssm_slots_per_req].copy_(ssm[:, heads]) + if source_mode == "replay": + backend = mem.linear.ssm_backend + backend.raw_keys.zero_() + backend.raw_values.zero_() + backend.gates.fill_(-0.5) + backend.betas.zero_() + backend.cursors[src_req] = 1 + ssm[:, heads].copy_(backend.get_accepted_state(src_req)) offset = 0 for state in mem.tails: end = offset + len(state.config.logical_layer_ids) @@ -173,6 +189,10 @@ def test_pd_kv_and_runtime_state_roundtrip( state.storage.zero_() mem.linear.conv_state.normal_() mem.linear.ssm_state.normal_() + if dest_mode == "replay": + mem.linear.ssm_backend.cursors[dst_req] = 3 + if mem.linear.mtp_state_index is not None: + mem.linear.mtp_state_index[dst_req] = decode_mtp_step for state in mem.tails: state.buffer.normal_() original_tails.append(torch.cat([state.buffer[:, dst_req] for state in mem.tails]).clone()) @@ -181,7 +201,7 @@ def test_pd_kv_and_runtime_state_roundtrip( x[:, torch.tensor([0, 1, 3], device="cuda") * stride].clone() for x, stride in ( (mem.linear.conv_state, 1), - (mem.linear.ssm_state, dst_mtp_size), + (mem.linear.ssm_state, dst_ssm_stride), *((state.buffer, 1) for state in mem.tails), ) ] @@ -235,18 +255,20 @@ def test_pd_kv_and_runtime_state_roundtrip( source_field.tensor[:, src_indexes].contiguous().view(torch.uint8), ) assert torch.equal(mem.linear.conv_state[:, dst_req, ..., :3], conv[:, :, heads].reshape(6, -1, 3)) - assert torch.equal(mem.linear.ssm_state[:, dst_req * dst_mtp_size], ssm[:, heads]) + assert torch.equal(mem.linear.ssm_state[:, dst_req * dst_ssm_stride], ssm[:, heads]) restored_tail = torch.cat([state.buffer[:, dst_req] for state in mem.tails]) expected_tail = original_tails[rank].clone() for position in range(length - 4, length): expected_tail[:, position % (4 + decode_mtp_step)] = tail[:, position % tail.shape[1]] assert torch.equal(restored_tail, expected_tail) + if dest_mode == "replay": + assert mem.linear.ssm_backend.cursors[dst_req] == 0 if mem.linear.mtp_state_index is not None: assert mem.linear.mtp_state_index[dst_req] == 0 for (state, stride), expected in zip( ( (mem.linear.conv_state, 1), - (mem.linear.ssm_state, dst_mtp_size), + (mem.linear.ssm_state, dst_ssm_stride), *((state.buffer, 1) for state in mem.tails), ), untouched[rank], From d3622dbba9216dc83d928444751ffb7fde9434ca Mon Sep 17 00:00:00 2001 From: sufubao Date: Tue, 3 Nov 2026 14:26:43 +0800 Subject: [PATCH 37/40] refactor: separate SSM state execution and acceptance --- docs/replay_ssm.md | 124 ++++ .../attention_state/bounded_att/base_state.py | 4 + .../bounded_att/conv_ssm_state.py | 556 ++++++++++++++++++ .../bounded_att/gdn/gdn_state.py | 490 +-------------- .../attention_state/bounded_att/interface.py | 9 +- .../bounded_att/kda/kda_state.py | 490 +-------------- .../sliding_window/fp8_e8m0_swa_state.py | 4 +- .../bounded_att/ssm/__init__.py | 34 ++ .../attention_state/bounded_att/ssm/base.py | 66 +++ .../bounded_att/ssm/compact.py | 177 ++++++ .../bounded_att/ssm/forward.py | 64 ++ .../attention_state/bounded_att/ssm/native.py | 96 +++ .../attention_state/bounded_att/ssm/replay.py | 300 ++++++++++ .../window_compressor_state.py | 4 +- lightllm/common/attention_state/commit.py | 59 ++ lightllm/common/attention_state/manager.py | 11 +- .../common/basemodel/attention/base_att.py | 13 +- .../common/basemodel/attention/linear/gdn.py | 213 ++----- .../common/basemodel/attention/linear/kda.py | 148 +---- .../basemodel/attention/linear/state.py | 89 +++ lightllm/common/basemodel/basemodel.py | 15 +- lightllm/common/basemodel/batch_objs.py | 4 +- lightllm/common/basemodel/cuda_graph.py | 4 +- .../triton_kernel/linear_att/compact_ssm.py | 156 ----- .../triton_kernel/linear_att/replayssm.py | 265 --------- lightllm/common/linear_state.py | 246 -------- lightllm/common/ssm_types.py | 13 + .../mode_backend/chunked_prefill/impl.py | 3 +- .../mode_backend/dp_backend/impl.py | 9 +- .../model_infer/mtp_speculative/utils.py | 30 +- test/benchmark/kernels/benchmark_kda_ssm.py | 10 +- .../test_hybrid_layer_mapping.py | 4 +- .../att_resources/test_mtp_state_commit.py | 30 +- .../basemodel/attention/linear/test_gdn.py | 31 +- .../basemodel/test_cuda_graph_autotune.py | 37 +- .../linear_att/test_acceptance_fusion.py | 52 +- .../linear_att/test_mtp_state_params.py | 22 +- .../linear_att/test_replayssm.py | 24 +- .../linear_att/test_replayssm_execution.py | 25 +- .../linear_att/test_replayssm_grouped_keys.py | 8 +- .../linear_att/test_replayssm_parallel.py | 10 +- .../linear_att/test_ssm_autotune.py | 35 +- .../linear_att/test_ssm_autotune_memory.py | 8 +- .../test_hybrid_req_manager_pd_sharing.py | 20 +- unit_tests/common/test_linear_state.py | 22 +- .../common/test_linear_state_contract.py | 237 +++++++- unit_tests/models/glm5_next/test_cache.py | 16 +- unit_tests/models/glm5_next/test_mtp.py | 6 +- unit_tests/models/glm5_next/test_pd_cache.py | 6 +- 49 files changed, 2140 insertions(+), 2159 deletions(-) create mode 100644 docs/replay_ssm.md create mode 100644 lightllm/common/attention_state/bounded_att/conv_ssm_state.py create mode 100644 lightllm/common/attention_state/bounded_att/ssm/__init__.py create mode 100644 lightllm/common/attention_state/bounded_att/ssm/base.py create mode 100644 lightllm/common/attention_state/bounded_att/ssm/compact.py create mode 100644 lightllm/common/attention_state/bounded_att/ssm/forward.py create mode 100644 lightllm/common/attention_state/bounded_att/ssm/native.py create mode 100644 lightllm/common/attention_state/bounded_att/ssm/replay.py create mode 100644 lightllm/common/attention_state/commit.py create mode 100644 lightllm/common/basemodel/attention/linear/state.py delete mode 100644 lightllm/common/linear_state.py create mode 100644 lightllm/common/ssm_types.py diff --git a/docs/replay_ssm.md b/docs/replay_ssm.md new file mode 100644 index 0000000000..c074a15ea5 --- /dev/null +++ b/docs/replay_ssm.md @@ -0,0 +1,124 @@ +# Recurrent state ownership + +GDN and KDA share the Conv/SSM resource lifecycle. Their attention implementations +own convolution, input packing and prefill computation. Each resource composes +one SSM implementation that owns recurrent execution, speculative storage and +accepted-state reconstruction. + +```mermaid +flowchart TD + Runner[Model runner: request layout and stream ordering] --> Attention[GDN / KDA attention] + Attention --> Forward[SSMForward: one resource and one invocation] + Forward --> Resource[ConvSSMState: accepted prefix and canonical snapshots] + Resource --> Native[NativeSSMState: candidate slots] + Resource --> Compact[CompactSSMState: verify scratch and reconstruction] + Resource --> Replay[ReplaySSMState: history and folding] + Native --> Kernels[Triton kernels] + Compact --> Kernels + Replay --> Kernels + Forward --> Output[ModelOutput.state_commits: immutable resource-keyed receipts] + Output --> Verify[MTP verification and logical microbatch partitions] + Verify --> Manager[AttentionStateManager: validate and dispatch] + Manager --> Resource + Resource --> Transfer[Canonical Conv/SSM pages: CPU cache / DP / PD] +``` + +## Lifetimes and interfaces + +| Owner | Lifetime | Responsibility | +|---|---|---| +| `ConvSSMState` | Worker/resource | Conv pool, accepted position, layer mapping, canonical save/restore | +| Private `SSMState` implementation | Worker/resource | SSM pool, candidate/history layout, kernels, acceptance | +| `SSMForward` | Forward or captured graph | Request/state indices, cumulative lengths, prior accepted counts, selected launch config | +| `StateCommit` | Model output or captured graph | Registered resource name and immutable implementation metadata | +| `StateCommitBatch` | Verification call | Number of real logical requests and receipts from their producing forward | + +The execution path is: + +1. `LinearDecodeAttState.init_state()` groups fixed or dynamic MTP rows once and + asks each participating resource to `prepare_decode()`. GDN obtains its actual + draft width from `MtpManager`; KDA retains its configured width. +2. GDN/KDA select the resource by role, algorithm and logical layer. After their + existing convolution/packing, `forward.run_layer()` dispatches the recurrence. + The attention caller does not inspect storage modes, cursors or kernel configs. +3. `forward.finish()` publishes a detached receipt when acceptance needs metadata. + Compact publishes the exact `SSMRunConfig` used for verification. Native and + Replay need no receipt. Multiple resources can publish independently. +4. `verify_mtp_tokens()` passes logical request IDs and accepted lengths to the + generic bounded-state commit interface. `MtpCommitContext` also retains the + expanded verification rows and partitions from the producing forwards. +5. The manager validates all receipts before dispatching any mutations. Each SSM + implementation commits its own state. Replay supplies its private cursor to + the existing fused accepted-position kernel; Compact reconstructs each + partition using its own forward's launch config. + +Logical request arrays are `[B]`; verification rows are `[T]`. Partition sizes +count real requests, excluding graph padding. Zero-sized partitions are allowed. +Registered resource names include the model role, so a draft receipt cannot be +applied to the main model. This uses the existing registration identity; it does +not add arbitrary same-type resource groups to the registration/transport schema. + +Missing Compact receipts, unknown resource names, duplicate receipts and invalid +partition totals are errors. No fallback reads a mutable "last forward" config. +The runner must commit exactly once, preserve request ordering, and prevent +overlapping live forwards from using the same real request slots. A receipt is +reusable across graph invocations; it is not a one-use transaction or rollback +mechanism. Tensor-value checks requiring a device synchronization are not added +to the decode path. + +## CUDA Graph and microbatch rules + +Request preparation runs before every invocation, including graph replay. +Replay snapshots its forward cursor before accepting the guaranteed token of +ordinary decode. Moving preparation solely into capture-time Python would break +later invocations. + +`SSMForward.copy_inputs()` copies nested tensor metadata into captured buffers, +validates their layouts, and preserves the captured launch config. Aliased +request/state-index buffers are copied once. Warmup clones each forward wrapper +so configuration selection cannot mutate the wrapper used for capture. Storage +and the tuning cache remain shared intentionally. + +Graph replay returns the captured `ModelOutput`; unpadding preserves its receipt +tuple. Overlapped microbatches pass separate `StateCommitBatch` entries, each with +its own real request count and configuration. Neither the scheduler nor the +generic manager needs to know that Compact uses those configs for reconstruction. + +## Checkpoints and numerical behavior + +`get_accepted_state()` returns the convolution window and SSM state at the same +accepted prefix. Native selects a candidate slot, Compact exposes its committed +state, and Replay materializes a snapshot without folding or changing live +history. CPU pages and DP/PD transport contain this canonical state, independent +of execution mode or speculative width. `prepare_restore()` clears history and +the accepted position before exposing destinations for both fields. + +Views may alias live storage. Existing stream ordering and completion events must +protect consuming copies before advancing or recycling requests. Restore callers +must fill both fields before execution resumes. A token length labels the +accepted prefix; it does not request reconstruction of arbitrary earlier tokens. + +Native and Compact retain per-token state rounding. Replay retains its FP32 +history accumulation and rounding at fold/materialization. These are intentionally +different numerical semantics. The refactor relocates host orchestration while +preserving the Triton recurrence/folding kernels and their launch parameters. +CLI modes, replay configuration, grouped heads, KDA gate bounds and existing +parallel-mode restrictions are unchanged. + +## Source map and verification + +- `common/attention_state/bounded_att/conv_ssm_state.py`: shared resource and + canonical page/transfer lifecycle; GDN/KDA subclasses supply identity/page type. +- `common/attention_state/bounded_att/ssm/`: native, compact, replay and forward + lifetimes; only the factory dispatches the configured mode. +- `common/basemodel/attention/linear/state.py`: shared request layout and graph + integration; `gdn.py`/`kda.py` retain algorithm-specific processing. +- `common/attention_state/commit.py`: generic verification and receipt contracts. +- `common/basemodel/triton_kernel/linear_att/`: numerical kernels and tuning. + +CPU contract tests cover accepted-state selection/reuse, resource-key routing, +invalid receipts, per-microbatch configs, nested graph copying and warmup +isolation. Kernel substitutions in these tests verify orchestration only. +The retained CUDA tests cover numerical recurrence, fused acceptance, dynamic +layouts, folding, graph alternation, checkpoint continuation and IPC. They must +run on compatible GPU hardware before claiming GPU or serving acceptance. diff --git a/lightllm/common/attention_state/bounded_att/base_state.py b/lightllm/common/attention_state/bounded_att/base_state.py index 6a96e76dc9..c852a5c7e9 100644 --- a/lightllm/common/attention_state/bounded_att/base_state.py +++ b/lightllm/common/attention_state/bounded_att/base_state.py @@ -7,3 +7,7 @@ class BoundedAttentionState(CoreInterface, CpuCacheInterface, DpCopyInterface, PdInterface): """单种 bounded-attention 方案的数据操作，索引和槽位由调用方分配。""" + + def validate_bounded_att_commit(self, context): + + pass diff --git a/lightllm/common/attention_state/bounded_att/conv_ssm_state.py b/lightllm/common/attention_state/bounded_att/conv_ssm_state.py new file mode 100644 index 0000000000..47f6e15515 --- /dev/null +++ b/lightllm/common/attention_state/bounded_att/conv_ssm_state.py @@ -0,0 +1,556 @@ +from __future__ import annotations + +from functools import lru_cache +from typing import TYPE_CHECKING, List, NamedTuple +from math import prod + +import torch + +from lightllm.utils.log_utils import init_logger + +from .ssm import create_ssm_state +from .ssm.forward import SSMForward + +from ..utils import deepcopy_result, dtype_name +from ..types import AxisMapping, StateTransferLayout, TensorField, TransferRegion +from .base_state import BoundedAttentionState +from .conv_ssm_types import ConvSSMStateConfig +from .transfer_utils import bounded_att_buffer_copies +from .copy_kernel import copy_bounded_att_tensors + +if TYPE_CHECKING: + from lightllm.common.attention_state.types import PageKind + + +logger = init_logger(__name__) + + +class ConvSSMTensors(NamedTuple): + conv: torch.Tensor + ssm: torch.Tensor + + +class ConvSSMState(BoundedAttentionState): + + algorithm: str + pages_pool_class: type + + def __init__( + self, + config: ConvSSMStateConfig, + request_capacity: int, + mtp_step: int = 0, + device: str | torch.device = "cuda", + *, + ssm_state_mode: str = "native", + ssm_replay_config: dict | None = None, + gate_lower_bound: float = -5.0, + ): + self.config = config + self.state_name = None + self.request_capacity = request_capacity + self.mtp_step = mtp_step + self._logical_to_storage_layer = {layer: index for index, layer in enumerate(config.logical_layer_ids)} + self.mtp_state_index = torch.zeros(request_capacity, dtype=torch.int32, device=device) if mtp_step else None + self.conv_state = torch.zeros( + (len(config.logical_layer_ids), request_capacity, config.conv_channels, config.conv_width + mtp_step), + dtype=config.conv_dtype, + device=device, + ) + self._ssm = create_ssm_state( + config, + request_capacity, + mtp_step, + device, + ssm_state_mode, + ssm_replay_config, + kda=self.algorithm == "kda", + gate_lower_bound=gate_lower_bound, + ) + self.big_pages_pool = None + self.small_pages_pool = None + + @property + def ssm_state(self): + return self._ssm.state + + @property + def ssm_slots_per_req(self): + return self._ssm.slots_per_request + + def _reset_tracking(self, req_idx): + self._ssm.clear_history(req_idx) + if self.mtp_state_index is not None: + self.mtp_state_index[req_idx].zero_() + + def reset(self, req_idx): + self._reset_tracking(req_idx) + self.conv_state[:, req_idx].zero_() + start = req_idx * self.ssm_slots_per_req + self.ssm_state[:, start : start + self.ssm_slots_per_req].zero_() + + def prepare_prefill(self, reqs): + self._ssm.merge_accepted_updates(reqs) + return reqs, reqs * self.ssm_slots_per_req + + def prepare_decode(self, reqs, cu_seqlens=None): + speculative = cu_seqlens is not None + accepted = None + if speculative: + if self.mtp_state_index is None: + raise ValueError("speculative decode requires MTP state slots") + accepted = self.mtp_state_index[reqs] + 1 + accepted = torch.where(reqs == self.request_capacity - 1, 1, accepted) + self._ssm.prepare_decode(reqs, cu_seqlens) + if not speculative: + self._ssm.accept_single_token(reqs) + return SSMForward(self, reqs, self._ssm.decode_indices(reqs, speculative), cu_seqlens, accepted) + + def get_accepted_state(self, req_idx): + + width = self.config.conv_width + accepted = self.mtp_state_index[req_idx] if self.mtp_state_index is not None else None + conv = self.conv_state[:, req_idx, ..., :width] + if accepted is not None: + offsets = torch.arange(width, device=conv.device) + accepted + conv = self.conv_state[:, req_idx].index_select(-1, offsets) + return ConvSSMTensors(conv, self._ssm.get_accepted_state(req_idx, accepted)) + + def prepare_restore(self, req_idx): + + self._reset_tracking(req_idx) + return ConvSSMTensors( + self.conv_state[:, req_idx, ..., : self.config.conv_width], + self.ssm_state[:, req_idx * self.ssm_slots_per_req], + ) + + def restore(self, req_idx, conv, ssm): + target = self.prepare_restore(req_idx) + target.conv.copy_(conv, non_blocking=True) + target.ssm.copy_(ssm, non_blocking=True) + + def get_storage_layer_index(self, logical_layer_id): + return self._logical_to_storage_layer[logical_layer_id] + + def get_state_type(self) -> str: + return ( + f"{self.algorithm}_conv_{dtype_name(dtype=self.config.conv_dtype)}" + f"_ssm_{dtype_name(dtype=self.config.ssm_dtype)}" + ) + + def initialize_page_pools( + self, + big_page_capacity: int, + small_page_capacity: int = 0, + device: str | torch.device | int = "cpu", + pin_memory: bool = True, + ) -> None: + """容量规划完成后创建本算法的大小页池；容量为零的页池保持 None。""" + assert self.big_pages_pool is None and self.small_pages_pool is None, "page pools already initialized" + assert big_page_capacity >= 0 and small_page_capacity >= 0, "page capacities must be nonnegative" + self.big_pages_pool = ( + self.pages_pool_class( + size=big_page_capacity, + config=self.config, + device=device, + pin_memory=pin_memory, + ) + if big_page_capacity + else None + ) + self.small_pages_pool = ( + self.pages_pool_class( + size=small_page_capacity, + config=self.config, + device=device, + pin_memory=pin_memory, + ) + if small_page_capacity + else None + ) + + # ======================================================================== + # 二、Bounded attention：计算视图、运行态重置、大小页与 MTP 提交 + # ======================================================================== + + def contains_logical_layer(self, logical_layer_id: int) -> bool: + """判断指定逻辑层是否属于当前 State 的层映射，共享层同样返回 True。""" + return logical_layer_id in self._logical_to_storage_layer + + def get_layer_view(self, logical_layer_id: int) -> tuple[torch.Tensor, torch.Tensor]: + """按逻辑层号映射到物理层，返回 Conv/SSM Tensor 引用；MTP 游标和步数直接从 State 读取。""" + storage_layer_index = self._logical_to_storage_layer[logical_layer_id] + return self.conv_state[storage_layer_index], self.ssm_state[storage_layer_index] + + def reset_bounded_att_state(self, req_indices: List[int]) -> None: + """清零指定请求的 Conv 历史、全部 SSM 候选及 MTP 游标，保留页池快照。""" + assert isinstance(req_indices, list), "request indices must be a list" + for req_idx in req_indices: + self.reset(req_idx) + + def save_bounded_att_runtime_to_pages( + self, + req_indices: List[int], + state_token_lens: List[int], + page_kind: PageKind, + page_slot_ids: List[int], + ) -> None: + """保存各请求指定绝对 token 长度处的基准状态；Conv/SSM 不依赖长度定位。""" + assert isinstance(req_indices, list) and isinstance( + page_slot_ids, list + ), "request and slot indices must be lists" + assert len(req_indices) == len(state_token_lens) == len(page_slot_ids), "request, length and page counts differ" + pool = self._pool(page_kind) + conv_copies, ssm_copies = [], [] + for req_idx, slot_id in zip(req_indices, page_slot_ids): + accepted = self.get_accepted_state(req_idx) + conv_copies.append((accepted.conv, pool.conv_state[slot_id])) + ssm_copies.append((accepted.ssm, pool.ssm_state[slot_id])) + copy_bounded_att_tensors(copies=conv_copies) + copy_bounded_att_tensors(copies=ssm_copies) + + def restore_bounded_att_runtime_from_pages( + self, + req_indices: List[int], + state_token_lens: List[int], + page_kind: PageKind, + page_slot_ids: List[int], + ) -> None: + """恢复各请求指定绝对 token 长度处的基准状态并归零游标；Conv/SSM 不依赖长度定位。""" + assert isinstance(req_indices, list) and isinstance( + page_slot_ids, list + ), "request and slot indices must be lists" + assert len(req_indices) == len(state_token_lens) == len(page_slot_ids), "request, length and page counts differ" + pool = self._pool(page_kind) + conv_copies, ssm_copies = [], [] + for req_idx, slot_id in zip(req_indices, page_slot_ids): + target = self.prepare_restore(req_idx) + conv_copies.append((pool.conv_state[slot_id], target.conv)) + ssm_copies.append((pool.ssm_state[slot_id], target.ssm)) + copy_bounded_att_tensors(copies=conv_copies) + copy_bounded_att_tensors(copies=ssm_copies) + + def validate_bounded_att_commit(self, context): + if self.mtp_state_index is not None: + self._ssm.validate_commit(context, self.state_name) + + def commit_bounded_att_state(self, req_indices, accepted_steps, *, context=None): + if context is not None: + context.validate(req_indices.numel()) + if self.mtp_state_index is not None: + self._ssm.commit(req_indices, accepted_steps, self.mtp_state_index, context, self.state_name) + + # ======================================================================== + # 三、CPU Cache 数据交换 + # ======================================================================== + + @deepcopy_result + @lru_cache(maxsize=None) + def get_cpu_transfer_regions(self) -> tuple[TransferRegion, ...]: + """返回已接受 Conv/SSM 快照的字段区间与切片，供 CPU 页布局和容量计算使用。 + + region.shape 描述 CPU 页中的全局字段；axis_mappings 将本地页池的 TP 分片映射到该字段。 + 起止位置相对本 State，由管理器平移到 CPU 联合页；不包含 MTP 候选区或游标。 + """ + return self._get_transfer_regions() + + def write_state_pages_to_cpu( + self, + page_kind: PageKind, + page_slot_ids: List[int], + cpu_page_buffers: torch.Tensor, + cpu_page_indices: List[int], + layout: StateTransferLayout, + cpu_page_ready_flags: list[bool] | None = None, + ) -> None: + """将指定页池槽位的快照写入 CPU 页。""" + assert layout.state_name == self.state_name, "transfer layout does not describe this resource" + assert 0 <= layout.byte_start < layout.byte_end, "invalid State byte interval" + assert cpu_page_buffers.device.type == "cpu", "CPU cache buffer must reside on CPU" + pool = self._pool(page_kind) + assert isinstance(page_slot_ids, list), "slot indices must be a list" + assert isinstance(cpu_page_indices, list), "CPU page indices must be a list" + assert len(page_slot_ids) == len(cpu_page_indices), "CPU page and slot counts differ" + if cpu_page_ready_flags is not None: + assert len(cpu_page_ready_flags) == len(cpu_page_indices), "CPU page and ready counts differ" + page_fields, page_indices = [], [] + for index, (slot, page) in enumerate(zip(page_slot_ids, cpu_page_indices)): + # ready 页不参与复制；同一传输区域的所有页面组成一个同形状批次。 + if cpu_page_ready_flags is not None and cpu_page_ready_flags[index]: + continue + page_fields.append(self._fields(pool.conv_state[slot], pool.ssm_state[slot])) + page_indices.append(page) + for copies in bounded_att_buffer_copies( + page_fields=page_fields, + page_buffers=cpu_page_buffers, + layout=layout, + page_indices=page_indices, + write=True, + ): + copy_bounded_att_tensors(copies=copies) + + def read_state_pages_from_cpu( + self, + cpu_page_buffers: torch.Tensor, + cpu_page_indices: List[int], + big_page_slot_ids: List[int], + layout: StateTransferLayout, + ) -> None: + """读入指定大页槽位；请求运行态另行通过 restore_bounded_att_runtime_from_pages 恢复。""" + assert layout.state_name == self.state_name, "transfer layout does not describe this resource" + assert 0 <= layout.byte_start < layout.byte_end, "invalid State byte interval" + assert cpu_page_buffers.device.type == "cpu", "CPU cache buffer must reside on CPU" + pool = self.big_pages_pool + assert isinstance(big_page_slot_ids, list), "slot indices must be a list" + assert isinstance(cpu_page_indices, list), "CPU page indices must be a list" + assert len(cpu_page_indices) == len(big_page_slot_ids), "CPU page and slot counts differ" + # 本地目标：选定大页槽位的 Conv/SSM 快照，保留本 TP rank 的分片形状。 + pool_page_fields = [self._fields(pool.conv_state[slot], pool.ssm_state[slot]) for slot in big_page_slot_ids] + # 外部来源：layout.regions 描述 CPU 全局字段和 TP 范围；write=False 将其读入本地页池。 + # 窗口保持快照顺序；请求 ring 的位置由后续 restore_bounded_att_runtime_from_pages 按长度恢复。 + for copies in bounded_att_buffer_copies( + page_fields=pool_page_fields, + page_buffers=cpu_page_buffers, + layout=layout, + page_indices=cpu_page_indices, + write=False, + ): + copy_bounded_att_tensors(copies=copies) + + # ======================================================================== + # 四、DP 迁移：大页快照与基准运行态 + # ======================================================================== + + def copy_big_pages_from_dp( + self, + src_page_ids: List[int], + dst_page_ids: List[int], + remote_dp_bounded_state: ConvSSMState, + ) -> None: + """复制同一 TP 分片的大页 Conv/SSM 快照，保留源页和双方请求运行态。""" + assert isinstance(src_page_ids, list) and isinstance(dst_page_ids, list), "page indices must be lists" + assert len(src_page_ids) == len(dst_page_ids), "source and destination page counts differ" + assert ( + type(remote_dp_bounded_state) is type(self) and remote_dp_bounded_state.config == self.config + ), "DP copy requires the same algorithm and equivalent TP state configurations" + src_pool = remote_dp_bounded_state.big_pages_pool + dst_pool = self.big_pages_pool + assert src_pool is not None and dst_pool is not None, "DP big page copy requires both page pools" + # 大页保存的是基准快照，与请求槽位及 MTP 候选区无关；页的引用和使用期限由调用方管理。 + # 页池均须在本进程注册；GPU 在当前流直接访问两端 CPU 存储，完成后才能归还页面。 + conv_copies, ssm_copies = [], [] + for src_page, dst_page in zip(src_page_ids, dst_page_ids): + conv_copies.append((src_pool.conv_state[src_page], dst_pool.conv_state[dst_page])) + ssm_copies.append((src_pool.ssm_state[src_page], dst_pool.ssm_state[dst_page])) + copy_bounded_att_tensors(copies=conv_copies) + copy_bounded_att_tensors(copies=ssm_copies) + + def copy_runtime_state_from_dp( + self, + src_req_indices: List[int], + dst_req_indices: List[int], + remote_dp_bounded_state: ConvSSMState, + ) -> None: + """复制同一 TP 分片的 prefill 基准运行态并重置目标游标，不写入页池。""" + assert isinstance(src_req_indices, list) and isinstance(dst_req_indices, list), "request indices must be lists" + assert len(src_req_indices) == len(dst_req_indices), "source and destination request counts differ" + assert ( + type(remote_dp_bounded_state) is type(self) and remote_dp_bounded_state.config == self.config + ), "DP copy requires the same algorithm and equivalent TP state configurations" + conv_copies, ssm_copies = [], [] + for src_req, dst_req in zip(src_req_indices, dst_req_indices): + source = remote_dp_bounded_state.get_accepted_state(src_req) + target = self.prepare_restore(dst_req) + conv_copies.append((source.conv, target.conv)) + ssm_copies.append((source.ssm, target.ssm)) + copy_bounded_att_tensors(copies=conv_copies) + copy_bounded_att_tensors(copies=ssm_copies) + + # ======================================================================== + # 五、PD 分离传输 + # ======================================================================== + + @deepcopy_result + @lru_cache(maxsize=None) + def get_pd_transfer_regions(self) -> tuple[TransferRegion, ...]: + """返回 PD 传输字段与切片，交由管理器安排其在联合页中的位置。""" + return self._get_transfer_regions() + + def write_to_pd_page( + self, + req_idx: int, + pd_page_buffers: torch.Tensor, + pd_page_index: int, + state_token_len: int, + layout: StateTransferLayout, + ) -> None: + """写出 state_token_len 处的已接受状态，不修改源 Replay 历史。""" + assert layout.state_name == self.state_name, "transfer layout does not describe this resource" + assert 0 <= layout.byte_start < layout.byte_end, "invalid State byte interval" + fields = self._fields(*self.get_accepted_state(req_idx)) + for copies in bounded_att_buffer_copies( + page_fields=[fields], + page_buffers=pd_page_buffers, + layout=layout, + page_indices=[pd_page_index], + write=True, + ): + copy_bounded_att_tensors(copies=copies) + + def read_from_pd_page( + self, + pd_page_buffers: torch.Tensor, + req_idx: int, + pd_page_index: int, + state_token_len: int, + layout: StateTransferLayout, + ) -> None: + """恢复 state_token_len 处的基准状态并规范化控制信息；Conv/SSM 不依赖长度定位。""" + assert layout.state_name == self.state_name, "transfer layout does not describe this resource" + assert 0 <= layout.byte_start < layout.byte_end, "invalid State byte interval" + fields = self._fields(*self.prepare_restore(req_idx)) + for copies in bounded_att_buffer_copies( + page_fields=[fields], + page_buffers=pd_page_buffers, + layout=layout, + page_indices=[pd_page_index], + write=False, + ): + copy_bounded_att_tensors(copies=copies) + if self.mtp_state_index is not None: + self.mtp_state_index[req_idx].fill_(0) + + # ======================================================================== + # 六、私有辅助实现 + # ======================================================================== + + def _get_transfer_regions(self) -> tuple[TransferRegion, ...]: + """依次描述全局 Conv 和 SSM 区域，并标明当前 TP rank 的读写范围。""" + # 阶段一：确定全局 Conv 形状；只传有效历史窗口，不包含 MTP 候选区和游标。 + config = self.config + storage_layer_count = len(config.logical_layer_ids) + conv_channels = 2 * config.global_k_heads * config.key_dim + config.global_v_heads * config.value_dim + conv_shape = (storage_layer_count, conv_channels, config.conv_width) + conv_byte_end = prod(conv_shape) * config.conv_dtype.itemsize + regions = [] + logger.debug( + "%s building transfer regions: TP=%d/%d, Conv shape=%s", + self.state_name, + config.tp_rank, + config.tp_world_size, + conv_shape, + ) + + # 阶段二：Conv 按 Q/K/V 拼接；三个映射共用同一块全局 Conv 区域，不能重复累计字节数。 + local_channel_start = 0 + global_channel_start = 0 + for name, local_head_count, global_head_count, head_dim in ( + ("Q", config.num_k_heads, config.global_k_heads, config.key_dim), + ("K", config.num_k_heads, config.global_k_heads, config.key_dim), + ("V", config.num_v_heads, config.global_v_heads, config.value_dim), + ): + # 配置已保证完整分片或副本。每组连续 rank 保存相同 heads，仅组内首 rank 写出。 + replica_count = config.tp_world_size * local_head_count // global_head_count + head_start = config.tp_rank // replica_count * local_head_count + write_enabled = config.tp_rank % replica_count == 0 + local_channel_count = local_head_count * head_dim + target_channel_start = global_channel_start + head_start * head_dim + conv_region = TransferRegion( + field="conv", + byte_start=0, + byte_end=conv_byte_end, + dtype=config.conv_dtype, + axes=("layer", "channel", "window"), + shape=conv_shape, + axis_mappings=( + AxisMapping( + axis="channel", + local_start=local_channel_start, + external_start=target_channel_start, + length=local_channel_count, + ), + AxisMapping( + axis="window", + local_start=0, + external_start=0, + length=config.conv_width, + ), + ), + write_enabled=write_enabled, + ) + regions.append(conv_region) + logger.debug( + "%s Conv %s: local_channel_start=%d, global_channel_start=%d, " + "channels=%d, replicas=%d, write_enabled=%s", + self.state_name, + name, + local_channel_start, + target_channel_start, + local_channel_count, + replica_count, + write_enabled, + ) + local_channel_start += local_channel_count + global_channel_start += global_head_count * head_dim + + # 阶段三：SSM 沿 V head 切分；接在 Conv 之后，按 SSM dtype 对齐字节起点。 + replica_count = config.tp_world_size * config.num_v_heads // config.global_v_heads + head_start = config.tp_rank // replica_count * config.num_v_heads + write_enabled = config.tp_rank % replica_count == 0 + ssm_shape = (storage_layer_count, config.global_v_heads, config.key_dim, config.value_dim) + element_bytes = config.ssm_dtype.itemsize + ssm_byte_start = (conv_byte_end + element_bytes - 1) // element_bytes * element_bytes + ssm_region = TransferRegion( + field="ssm", + byte_start=ssm_byte_start, + byte_end=ssm_byte_start + prod(ssm_shape) * element_bytes, + dtype=config.ssm_dtype, + axes=("layer", "head", "key_dim", "value_dim"), + shape=ssm_shape, + axis_mappings=( + AxisMapping( + axis="head", + local_start=0, + external_start=head_start, + length=config.num_v_heads, + ), + AxisMapping( + axis="key_dim", + local_start=0, + external_start=0, + length=config.key_dim, + ), + AxisMapping( + axis="value_dim", + local_start=0, + external_start=0, + length=config.value_dim, + ), + ), + write_enabled=write_enabled, + ) + regions.append(ssm_region) + logger.debug( + "%s transfer regions built: SSM head_start=%d, heads=%d, write_enabled=%s, SSM bytes=[%d, %d)", + self.state_name, + head_start, + config.num_v_heads, + write_enabled, + ssm_region.byte_start, + ssm_region.byte_end, + ) + return tuple(regions) + + def _pool(self, page_kind: PageKind): + """根据大小页类型取得已启用的页池，不访问 allocator。""" + assert page_kind in ("big", "small"), f"unknown page kind {page_kind}" + pool = self.big_pages_pool if page_kind == "big" else self.small_pages_pool + assert pool is not None, f"{page_kind} pages are disabled for {self.state_name}" + return pool + + @staticmethod + def _fields(conv: torch.Tensor, ssm: torch.Tensor): + return { + "conv": TensorField(tensor=conv, axes=("layer", "channel", "window")), + "ssm": TensorField(tensor=ssm, axes=("layer", "head", "key_dim", "value_dim")), + } diff --git a/lightllm/common/attention_state/bounded_att/gdn/gdn_state.py b/lightllm/common/attention_state/bounded_att/gdn/gdn_state.py index fd31143063..59c64b8ea0 100644 --- a/lightllm/common/attention_state/bounded_att/gdn/gdn_state.py +++ b/lightllm/common/attention_state/bounded_att/gdn/gdn_state.py @@ -1,490 +1,8 @@ -from __future__ import annotations - -from functools import lru_cache -from typing import TYPE_CHECKING, List -from math import prod - -import torch - -from lightllm.utils.log_utils import init_logger -from lightllm.common.linear_state import LinearStateManager - -from ...utils import deepcopy_result, dtype_name -from ...types import AxisMapping, StateTransferLayout, TensorField, TransferRegion -from ..base_state import BoundedAttentionState -from ..conv_ssm_types import ConvSSMStateConfig +from ..conv_ssm_state import ConvSSMState from .gdn_pages_pool import GDNPagesPool -from ..transfer_utils import bounded_att_buffer_copies -from ..copy_kernel import copy_bounded_att_tensors - -if TYPE_CHECKING: - from lightllm.common.attention_state.types import PageKind - - -logger = init_logger(__name__) - - -class GDNState(LinearStateManager, BoundedAttentionState): - """Gated DeltaNet 的 Conv 历史、SSM 和 MTP 游标，支持不同的 K/V head 数及维度。""" - - big_pages_pool: GDNPagesPool | None - small_pages_pool: GDNPagesPool | None - - # ======================================================================== - # 一、资源定义与初始化 - # ======================================================================== - - def __init__( - self, - config: ConvSSMStateConfig, - request_capacity: int, - mtp_step: int = 0, - device: str | torch.device = "cuda", - *, - ssm_state_mode: str = "native", - ssm_replay_config: dict | None = None, - gate_lower_bound: float = -5.0, - ) -> None: - """按请求容量创建 Conv/SSM 运行态和 MTP 游标；大小页池延后初始化。""" - assert isinstance(config, ConvSSMStateConfig), "config must be ConvSSMStateConfig" - self.config = config - self.state_name: str | None = None # 由 AttentionStateManager.register_state() 绑定时设置。 - LinearStateManager.__init__( - self, - config, - request_capacity, - mtp_step, - device, - ssm_state_mode=ssm_state_mode, - ssm_replay_config=ssm_replay_config, - kda=False, - gate_lower_bound=gate_lower_bound, - ) - self.big_pages_pool = None - self.small_pages_pool = None - - def get_state_type(self) -> str: - """返回 GDN 算法及 Conv/SSM 存储精度。""" - return f"gdn_conv_{dtype_name(dtype=self.config.conv_dtype)}" f"_ssm_{dtype_name(dtype=self.config.ssm_dtype)}" - - def initialize_page_pools( - self, - big_page_capacity: int, - small_page_capacity: int = 0, - device: str | torch.device | int = "cpu", - pin_memory: bool = True, - ) -> None: - """容量规划完成后创建本算法的大小页池；容量为零的页池保持 None。""" - assert self.big_pages_pool is None and self.small_pages_pool is None, "page pools already initialized" - assert big_page_capacity >= 0 and small_page_capacity >= 0, "page capacities must be nonnegative" - self.big_pages_pool = ( - GDNPagesPool( - size=big_page_capacity, - config=self.config, - device=device, - pin_memory=pin_memory, - ) - if big_page_capacity - else None - ) - self.small_pages_pool = ( - GDNPagesPool( - size=small_page_capacity, - config=self.config, - device=device, - pin_memory=pin_memory, - ) - if small_page_capacity - else None - ) - - # ======================================================================== - # 二、Bounded attention：计算视图、运行态重置、大小页与 MTP 提交 - # ======================================================================== - - def contains_logical_layer(self, logical_layer_id: int) -> bool: - """判断指定逻辑层是否属于当前 State 的层映射，共享层同样返回 True。""" - return logical_layer_id in self._logical_to_storage_layer - - def get_layer_view(self, logical_layer_id: int) -> tuple[torch.Tensor, torch.Tensor]: - """按逻辑层号映射到物理层，返回 Conv/SSM Tensor 引用；MTP 游标和步数直接从 State 读取。""" - storage_layer_index = self._logical_to_storage_layer[logical_layer_id] - return self.conv_state[storage_layer_index], self.ssm_state[storage_layer_index] - - def reset_bounded_att_state(self, req_indices: List[int]) -> None: - """清零指定请求的 Conv 历史、全部 SSM 候选及 MTP 游标，保留页池快照。""" - assert isinstance(req_indices, list), "request indices must be a list" - for req_idx in req_indices: - self.reset(req_idx) - - def save_bounded_att_runtime_to_pages( - self, - req_indices: List[int], - state_token_lens: List[int], - page_kind: PageKind, - page_slot_ids: List[int], - ) -> None: - """保存各请求指定绝对 token 长度处的基准状态；Conv/SSM 不依赖长度定位。""" - assert isinstance(req_indices, list) and isinstance( - page_slot_ids, list - ), "request and slot indices must be lists" - assert len(req_indices) == len(state_token_lens) == len(page_slot_ids), "request, length and page counts differ" - pool = self._pool(page_kind) - conv_copies, ssm_copies = [], [] - for req_idx, slot_id in zip(req_indices, page_slot_ids): - accepted = self.get_accepted_state(req_idx) - conv_copies.append((accepted.conv, pool.conv_state[slot_id])) - ssm_copies.append((accepted.ssm, pool.ssm_state[slot_id])) - copy_bounded_att_tensors(copies=conv_copies) - copy_bounded_att_tensors(copies=ssm_copies) - - def restore_bounded_att_runtime_from_pages( - self, - req_indices: List[int], - state_token_lens: List[int], - page_kind: PageKind, - page_slot_ids: List[int], - ) -> None: - """恢复各请求指定绝对 token 长度处的基准状态并归零游标；Conv/SSM 不依赖长度定位。""" - assert isinstance(req_indices, list) and isinstance( - page_slot_ids, list - ), "request and slot indices must be lists" - assert len(req_indices) == len(state_token_lens) == len(page_slot_ids), "request, length and page counts differ" - pool = self._pool(page_kind) - conv_copies, ssm_copies = [], [] - for req_idx, slot_id in zip(req_indices, page_slot_ids): - target = self.prepare_restore(req_idx) - conv_copies.append((pool.conv_state[slot_id], target.conv)) - ssm_copies.append((pool.ssm_state[slot_id], target.ssm)) - copy_bounded_att_tensors(copies=conv_copies) - copy_bounded_att_tensors(copies=ssm_copies) - if self.mtp_state_index is not None: - for req_idx in req_indices: - self.mtp_state_index[req_idx].fill_(0) - - def commit_bounded_att_state( - self, - req_indices: torch.Tensor, - accepted_steps: torch.Tensor, - ) -> None: - """用 CUDA 请求索引和 int32 接受步数选择 SSM 候选位置；步数包含基准输入步。""" - assert req_indices.is_cuda and accepted_steps.is_cuda, "commit inputs must be CUDA tensors" - if self.mtp_state_index is not None: - self.mtp_state_index[req_indices] = accepted_steps - 1 - if self.ssm_backend is not None: - self.ssm_backend.accept_updates(req_indices, self.mtp_state_index) - - # ======================================================================== - # 三、CPU Cache 数据交换 - # ======================================================================== - - @deepcopy_result - @lru_cache(maxsize=None) - def get_cpu_transfer_regions(self) -> tuple[TransferRegion, ...]: - """返回已接受 Conv/SSM 快照的字段区间与切片，供 CPU 页布局和容量计算使用。 - - region.shape 描述 CPU 页中的全局字段；axis_mappings 将本地页池的 TP 分片映射到该字段。 - 起止位置相对本 State，由管理器平移到 CPU 联合页；不包含 MTP 候选区或游标。 - """ - return self._get_transfer_regions() - - def write_state_pages_to_cpu( - self, - page_kind: PageKind, - page_slot_ids: List[int], - cpu_page_buffers: torch.Tensor, - cpu_page_indices: List[int], - layout: StateTransferLayout, - cpu_page_ready_flags: list[bool] | None = None, - ) -> None: - """将指定页池槽位的快照写入 CPU 页。""" - assert layout.state_name == self.state_name, "transfer layout does not describe this resource" - assert 0 <= layout.byte_start < layout.byte_end, "invalid State byte interval" - assert cpu_page_buffers.device.type == "cpu", "CPU cache buffer must reside on CPU" - pool = self._pool(page_kind) - assert isinstance(page_slot_ids, list), "slot indices must be a list" - assert isinstance(cpu_page_indices, list), "CPU page indices must be a list" - assert len(page_slot_ids) == len(cpu_page_indices), "CPU page and slot counts differ" - if cpu_page_ready_flags is not None: - assert len(cpu_page_ready_flags) == len(cpu_page_indices), "CPU page and ready counts differ" - page_fields, page_indices = [], [] - for index, (slot, page) in enumerate(zip(page_slot_ids, cpu_page_indices)): - # ready 页不参与复制；同一传输区域的所有页面组成一个同形状批次。 - if cpu_page_ready_flags is not None and cpu_page_ready_flags[index]: - continue - page_fields.append(self._fields(pool.conv_state[slot], pool.ssm_state[slot])) - page_indices.append(page) - for copies in bounded_att_buffer_copies( - page_fields=page_fields, - page_buffers=cpu_page_buffers, - layout=layout, - page_indices=page_indices, - write=True, - ): - copy_bounded_att_tensors(copies=copies) - - def read_state_pages_from_cpu( - self, - cpu_page_buffers: torch.Tensor, - cpu_page_indices: List[int], - big_page_slot_ids: List[int], - layout: StateTransferLayout, - ) -> None: - """读入指定大页槽位；请求运行态另行通过 restore_bounded_att_runtime_from_pages 恢复。""" - assert layout.state_name == self.state_name, "transfer layout does not describe this resource" - assert 0 <= layout.byte_start < layout.byte_end, "invalid State byte interval" - assert cpu_page_buffers.device.type == "cpu", "CPU cache buffer must reside on CPU" - pool = self.big_pages_pool - assert isinstance(big_page_slot_ids, list), "slot indices must be a list" - assert isinstance(cpu_page_indices, list), "CPU page indices must be a list" - assert len(cpu_page_indices) == len(big_page_slot_ids), "CPU page and slot counts differ" - # 本地目标：选定大页槽位的 Conv/SSM 快照，保留本 TP rank 的分片形状。 - pool_page_fields = [self._fields(pool.conv_state[slot], pool.ssm_state[slot]) for slot in big_page_slot_ids] - # 外部来源：layout.regions 描述 CPU 全局字段和 TP 范围；write=False 将其读入本地页池。 - # 窗口保持快照顺序；请求 ring 的位置由后续 restore_bounded_att_runtime_from_pages 按长度恢复。 - for copies in bounded_att_buffer_copies( - page_fields=pool_page_fields, - page_buffers=cpu_page_buffers, - layout=layout, - page_indices=cpu_page_indices, - write=False, - ): - copy_bounded_att_tensors(copies=copies) - - # ======================================================================== - # 四、DP 迁移：大页快照与基准运行态 - # ======================================================================== - - def copy_big_pages_from_dp( - self, - src_page_ids: List[int], - dst_page_ids: List[int], - remote_dp_bounded_state: GDNState, - ) -> None: - """复制同一 TP 分片的大页 Conv/SSM 快照，保留源页和双方请求运行态。""" - assert isinstance(src_page_ids, list) and isinstance(dst_page_ids, list), "page indices must be lists" - assert len(src_page_ids) == len(dst_page_ids), "source and destination page counts differ" - assert ( - type(remote_dp_bounded_state) is type(self) and remote_dp_bounded_state.config == self.config - ), "DP copy requires the same algorithm and equivalent TP state configurations" - src_pool = remote_dp_bounded_state.big_pages_pool - dst_pool = self.big_pages_pool - assert src_pool is not None and dst_pool is not None, "DP big page copy requires both page pools" - # 大页保存的是基准快照，与请求槽位及 MTP 候选区无关；页的引用和使用期限由调用方管理。 - # 页池均须在本进程注册；GPU 在当前流直接访问两端 CPU 存储，完成后才能归还页面。 - conv_copies, ssm_copies = [], [] - for src_page, dst_page in zip(src_page_ids, dst_page_ids): - conv_copies.append((src_pool.conv_state[src_page], dst_pool.conv_state[dst_page])) - ssm_copies.append((src_pool.ssm_state[src_page], dst_pool.ssm_state[dst_page])) - copy_bounded_att_tensors(copies=conv_copies) - copy_bounded_att_tensors(copies=ssm_copies) - - def copy_runtime_state_from_dp( - self, - src_req_indices: List[int], - dst_req_indices: List[int], - remote_dp_bounded_state: GDNState, - ) -> None: - """复制同一 TP 分片的 prefill 基准运行态并重置目标游标，不写入页池。""" - assert isinstance(src_req_indices, list) and isinstance(dst_req_indices, list), "request indices must be lists" - assert len(src_req_indices) == len(dst_req_indices), "source and destination request counts differ" - assert ( - type(remote_dp_bounded_state) is type(self) and remote_dp_bounded_state.config == self.config - ), "DP copy requires the same algorithm and equivalent TP state configurations" - conv_copies, ssm_copies = [], [] - for src_req, dst_req in zip(src_req_indices, dst_req_indices): - source = remote_dp_bounded_state.get_accepted_state(src_req) - target = self.prepare_restore(dst_req) - conv_copies.append((source.conv, target.conv)) - ssm_copies.append((source.ssm, target.ssm)) - copy_bounded_att_tensors(copies=conv_copies) - copy_bounded_att_tensors(copies=ssm_copies) - - # ======================================================================== - # 五、PD 分离传输 - # ======================================================================== - - @deepcopy_result - @lru_cache(maxsize=None) - def get_pd_transfer_regions(self) -> tuple[TransferRegion, ...]: - """返回 PD 传输字段与切片，交由管理器安排其在联合页中的位置。""" - return self._get_transfer_regions() - - def write_to_pd_page( - self, - req_idx: int, - pd_page_buffers: torch.Tensor, - pd_page_index: int, - state_token_len: int, - layout: StateTransferLayout, - ) -> None: - """写出 state_token_len 处的已接受状态，不修改源 Replay 历史。""" - assert layout.state_name == self.state_name, "transfer layout does not describe this resource" - assert 0 <= layout.byte_start < layout.byte_end, "invalid State byte interval" - fields = self._fields(*self.get_accepted_state(req_idx)) - for copies in bounded_att_buffer_copies( - page_fields=[fields], - page_buffers=pd_page_buffers, - layout=layout, - page_indices=[pd_page_index], - write=True, - ): - copy_bounded_att_tensors(copies=copies) - - def read_from_pd_page( - self, - pd_page_buffers: torch.Tensor, - req_idx: int, - pd_page_index: int, - state_token_len: int, - layout: StateTransferLayout, - ) -> None: - """恢复 state_token_len 处的基准状态并规范化控制信息；Conv/SSM 不依赖长度定位。""" - assert layout.state_name == self.state_name, "transfer layout does not describe this resource" - assert 0 <= layout.byte_start < layout.byte_end, "invalid State byte interval" - fields = self._fields(*self.prepare_restore(req_idx)) - for copies in bounded_att_buffer_copies( - page_fields=[fields], - page_buffers=pd_page_buffers, - layout=layout, - page_indices=[pd_page_index], - write=False, - ): - copy_bounded_att_tensors(copies=copies) - if self.mtp_state_index is not None: - self.mtp_state_index[req_idx].fill_(0) - - # ======================================================================== - # 六、私有辅助实现 - # ======================================================================== - - def _get_transfer_regions(self) -> tuple[TransferRegion, ...]: - """依次描述全局 Conv 和 SSM 区域，并标明当前 TP rank 的读写范围。""" - # 阶段一：确定全局 Conv 形状；只传有效历史窗口，不包含 MTP 候选区和游标。 - config = self.config - storage_layer_count = len(config.logical_layer_ids) - conv_channels = 2 * config.global_k_heads * config.key_dim + config.global_v_heads * config.value_dim - conv_shape = (storage_layer_count, conv_channels, config.conv_width) - conv_byte_end = prod(conv_shape) * config.conv_dtype.itemsize - regions = [] - logger.debug( - "%s building transfer regions: TP=%d/%d, Conv shape=%s", - self.state_name, - config.tp_rank, - config.tp_world_size, - conv_shape, - ) - - # 阶段二：Conv 按 Q/K/V 拼接；三个映射共用同一块全局 Conv 区域，不能重复累计字节数。 - local_channel_start = 0 - global_channel_start = 0 - for name, local_head_count, global_head_count, head_dim in ( - ("Q", config.num_k_heads, config.global_k_heads, config.key_dim), - ("K", config.num_k_heads, config.global_k_heads, config.key_dim), - ("V", config.num_v_heads, config.global_v_heads, config.value_dim), - ): - # 配置已保证完整分片或副本。每组连续 rank 保存相同 heads，仅组内首 rank 写出。 - replica_count = config.tp_world_size * local_head_count // global_head_count - head_start = config.tp_rank // replica_count * local_head_count - write_enabled = config.tp_rank % replica_count == 0 - local_channel_count = local_head_count * head_dim - target_channel_start = global_channel_start + head_start * head_dim - conv_region = TransferRegion( - field="conv", - byte_start=0, - byte_end=conv_byte_end, - dtype=config.conv_dtype, - axes=("layer", "channel", "window"), - shape=conv_shape, - axis_mappings=( - AxisMapping( - axis="channel", - local_start=local_channel_start, - external_start=target_channel_start, - length=local_channel_count, - ), - AxisMapping( - axis="window", - local_start=0, - external_start=0, - length=config.conv_width, - ), - ), - write_enabled=write_enabled, - ) - regions.append(conv_region) - logger.debug( - "%s Conv %s: local_channel_start=%d, global_channel_start=%d, " - "channels=%d, replicas=%d, write_enabled=%s", - self.state_name, - name, - local_channel_start, - target_channel_start, - local_channel_count, - replica_count, - write_enabled, - ) - local_channel_start += local_channel_count - global_channel_start += global_head_count * head_dim - # 阶段三：SSM 沿 V head 切分；接在 Conv 之后，按 SSM dtype 对齐字节起点。 - replica_count = config.tp_world_size * config.num_v_heads // config.global_v_heads - head_start = config.tp_rank // replica_count * config.num_v_heads - write_enabled = config.tp_rank % replica_count == 0 - ssm_shape = (storage_layer_count, config.global_v_heads, config.key_dim, config.value_dim) - element_bytes = config.ssm_dtype.itemsize - ssm_byte_start = (conv_byte_end + element_bytes - 1) // element_bytes * element_bytes - ssm_region = TransferRegion( - field="ssm", - byte_start=ssm_byte_start, - byte_end=ssm_byte_start + prod(ssm_shape) * element_bytes, - dtype=config.ssm_dtype, - axes=("layer", "head", "key_dim", "value_dim"), - shape=ssm_shape, - axis_mappings=( - AxisMapping( - axis="head", - local_start=0, - external_start=head_start, - length=config.num_v_heads, - ), - AxisMapping( - axis="key_dim", - local_start=0, - external_start=0, - length=config.key_dim, - ), - AxisMapping( - axis="value_dim", - local_start=0, - external_start=0, - length=config.value_dim, - ), - ), - write_enabled=write_enabled, - ) - regions.append(ssm_region) - logger.debug( - "%s transfer regions built: SSM head_start=%d, heads=%d, write_enabled=%s, SSM bytes=[%d, %d)", - self.state_name, - head_start, - config.num_v_heads, - write_enabled, - ssm_region.byte_start, - ssm_region.byte_end, - ) - return tuple(regions) - def _pool(self, page_kind: PageKind): - """根据大小页类型取得已启用的页池，不访问 allocator。""" - assert page_kind in ("big", "small"), f"unknown page kind {page_kind}" - pool = self.big_pages_pool if page_kind == "big" else self.small_pages_pool - assert pool is not None, f"{page_kind} pages are disabled for {self.state_name}" - return pool +class GDNState(ConvSSMState): - @staticmethod - def _fields(conv: torch.Tensor, ssm: torch.Tensor): - return { - "conv": TensorField(tensor=conv, axes=("layer", "channel", "window")), - "ssm": TensorField(tensor=ssm, axes=("layer", "head", "key_dim", "value_dim")), - } + algorithm = "gdn" + pages_pool_class = GDNPagesPool diff --git a/lightllm/common/attention_state/bounded_att/interface.py b/lightllm/common/attention_state/bounded_att/interface.py index fe1fceb3bd..fa5d1e01ee 100644 --- a/lightllm/common/attention_state/bounded_att/interface.py +++ b/lightllm/common/attention_state/bounded_att/interface.py @@ -102,8 +102,9 @@ def commit_bounded_att_state( self, req_indices: torch.Tensor, accepted_steps: torch.Tensor, + *, + context=None, ) -> None: - """两个输入均为 CUDA Tensor；int32 接受步数包含基准输入步，候选区与游标由实现管理。""" ... @@ -173,10 +174,7 @@ def copy_runtime_state_from_dp( dst_req_indices: List[int], remote_dp_bounded_state: CoreInterface, ) -> None: - """复制其他 DP 同一 TP 分片的 prefill 基准运行态并重置目标游标，不修改页池。 - 调用方保证源状态尚未执行 MTP 提交；KDA/GDN 的源游标必须为 0。 - """ ... @@ -205,10 +203,7 @@ def write_to_pd_page( state_token_len: int, layout: StateTransferLayout, ) -> None: - """将指定槽位在 state_token_len 绝对长度处的 prefill 基准运行态写入 PD 页。 - 调用方保证源状态尚未执行 MTP 提交；KDA/GDN 的源游标必须为 0。 - """ ... @abstractmethod diff --git a/lightllm/common/attention_state/bounded_att/kda/kda_state.py b/lightllm/common/attention_state/bounded_att/kda/kda_state.py index 9f8051836f..903da44b68 100644 --- a/lightllm/common/attention_state/bounded_att/kda/kda_state.py +++ b/lightllm/common/attention_state/bounded_att/kda/kda_state.py @@ -1,490 +1,8 @@ -from __future__ import annotations - -from functools import lru_cache -from typing import TYPE_CHECKING, List -from math import prod - -import torch - -from lightllm.utils.log_utils import init_logger -from lightllm.common.linear_state import LinearStateManager - -from ...utils import deepcopy_result, dtype_name -from ...types import AxisMapping, StateTransferLayout, TensorField, TransferRegion -from ..base_state import BoundedAttentionState -from ..conv_ssm_types import ConvSSMStateConfig +from ..conv_ssm_state import ConvSSMState from .kda_pages_pool import KDAPagesPool -from ..transfer_utils import bounded_att_buffer_copies -from ..copy_kernel import copy_bounded_att_tensors - -if TYPE_CHECKING: - from lightllm.common.attention_state.types import PageKind - - -logger = init_logger(__name__) - - -class KDAState(LinearStateManager, BoundedAttentionState): - """KDA 的 Conv 历史、方形 SSM 和 MTP 游标，要求 Q/K/V head 数及维度相同。""" - - big_pages_pool: KDAPagesPool | None - small_pages_pool: KDAPagesPool | None - - # ======================================================================== - # 一、资源定义与初始化 - # ======================================================================== - - def __init__( - self, - config: ConvSSMStateConfig, - request_capacity: int, - mtp_step: int = 0, - device: str | torch.device = "cuda", - *, - ssm_state_mode: str = "native", - ssm_replay_config: dict | None = None, - gate_lower_bound: float = -5.0, - ) -> None: - """按请求容量创建 Conv/SSM 运行态和 MTP 游标；大小页池延后初始化。""" - assert isinstance(config, ConvSSMStateConfig), "config must be ConvSSMStateConfig" - self.config = config - self.state_name: str | None = None # 由 AttentionStateManager.register_state() 绑定时设置。 - LinearStateManager.__init__( - self, - config, - request_capacity, - mtp_step, - device, - ssm_state_mode=ssm_state_mode, - ssm_replay_config=ssm_replay_config, - kda=True, - gate_lower_bound=gate_lower_bound, - ) - self.big_pages_pool = None - self.small_pages_pool = None - - def get_state_type(self) -> str: - """返回 KDA 算法及 Conv/SSM 存储精度。""" - return f"kda_conv_{dtype_name(dtype=self.config.conv_dtype)}" f"_ssm_{dtype_name(dtype=self.config.ssm_dtype)}" - - def initialize_page_pools( - self, - big_page_capacity: int, - small_page_capacity: int = 0, - device: str | torch.device | int = "cpu", - pin_memory: bool = True, - ) -> None: - """容量规划完成后创建本算法的大小页池；容量为零的页池保持 None。""" - assert self.big_pages_pool is None and self.small_pages_pool is None, "page pools already initialized" - assert big_page_capacity >= 0 and small_page_capacity >= 0, "page capacities must be nonnegative" - self.big_pages_pool = ( - KDAPagesPool( - size=big_page_capacity, - config=self.config, - device=device, - pin_memory=pin_memory, - ) - if big_page_capacity - else None - ) - self.small_pages_pool = ( - KDAPagesPool( - size=small_page_capacity, - config=self.config, - device=device, - pin_memory=pin_memory, - ) - if small_page_capacity - else None - ) - - # ======================================================================== - # 二、Bounded attention：计算视图、运行态重置、大小页与 MTP 提交 - # ======================================================================== - - def contains_logical_layer(self, logical_layer_id: int) -> bool: - """判断指定逻辑层是否属于当前 State 的层映射，共享层同样返回 True。""" - return logical_layer_id in self._logical_to_storage_layer - - def get_layer_view(self, logical_layer_id: int) -> tuple[torch.Tensor, torch.Tensor]: - """按逻辑层号映射到物理层，返回 Conv/SSM Tensor 引用；MTP 游标和步数直接从 State 读取。""" - storage_layer_index = self._logical_to_storage_layer[logical_layer_id] - return self.conv_state[storage_layer_index], self.ssm_state[storage_layer_index] - - def reset_bounded_att_state(self, req_indices: List[int]) -> None: - """清零指定请求的 Conv 历史、全部 SSM 候选及 MTP 游标，保留页池快照。""" - assert isinstance(req_indices, list), "request indices must be a list" - for req_idx in req_indices: - self.reset(req_idx) - - def save_bounded_att_runtime_to_pages( - self, - req_indices: List[int], - state_token_lens: List[int], - page_kind: PageKind, - page_slot_ids: List[int], - ) -> None: - """保存各请求指定绝对 token 长度处的基准状态；Conv/SSM 不依赖长度定位。""" - assert isinstance(req_indices, list) and isinstance( - page_slot_ids, list - ), "request and slot indices must be lists" - assert len(req_indices) == len(state_token_lens) == len(page_slot_ids), "request, length and page counts differ" - pool = self._pool(page_kind) - conv_copies, ssm_copies = [], [] - for req_idx, slot_id in zip(req_indices, page_slot_ids): - accepted = self.get_accepted_state(req_idx) - conv_copies.append((accepted.conv, pool.conv_state[slot_id])) - ssm_copies.append((accepted.ssm, pool.ssm_state[slot_id])) - copy_bounded_att_tensors(copies=conv_copies) - copy_bounded_att_tensors(copies=ssm_copies) - - def restore_bounded_att_runtime_from_pages( - self, - req_indices: List[int], - state_token_lens: List[int], - page_kind: PageKind, - page_slot_ids: List[int], - ) -> None: - """恢复各请求指定绝对 token 长度处的基准状态并归零游标；Conv/SSM 不依赖长度定位。""" - assert isinstance(req_indices, list) and isinstance( - page_slot_ids, list - ), "request and slot indices must be lists" - assert len(req_indices) == len(state_token_lens) == len(page_slot_ids), "request, length and page counts differ" - pool = self._pool(page_kind) - conv_copies, ssm_copies = [], [] - for req_idx, slot_id in zip(req_indices, page_slot_ids): - target = self.prepare_restore(req_idx) - conv_copies.append((pool.conv_state[slot_id], target.conv)) - ssm_copies.append((pool.ssm_state[slot_id], target.ssm)) - copy_bounded_att_tensors(copies=conv_copies) - copy_bounded_att_tensors(copies=ssm_copies) - if self.mtp_state_index is not None: - for req_idx in req_indices: - self.mtp_state_index[req_idx].fill_(0) - - def commit_bounded_att_state( - self, - req_indices: torch.Tensor, - accepted_steps: torch.Tensor, - ) -> None: - """用 CUDA 请求索引和 int32 接受步数选择 SSM 候选位置；步数包含基准输入步。""" - assert req_indices.is_cuda and accepted_steps.is_cuda, "commit inputs must be CUDA tensors" - if self.mtp_state_index is not None: - self.mtp_state_index[req_indices] = accepted_steps - 1 - if self.ssm_backend is not None: - self.ssm_backend.accept_updates(req_indices, self.mtp_state_index) - - # ======================================================================== - # 三、CPU Cache 数据交换 - # ======================================================================== - - @deepcopy_result - @lru_cache(maxsize=None) - def get_cpu_transfer_regions(self) -> tuple[TransferRegion, ...]: - """返回已接受 Conv/SSM 快照的字段区间与切片，供 CPU 页布局和容量计算使用。 - - region.shape 描述 CPU 页中的全局字段；axis_mappings 将本地页池的 TP 分片映射到该字段。 - 起止位置相对本 State，由管理器平移到 CPU 联合页；不包含 MTP 候选区或游标。 - """ - return self._get_transfer_regions() - - def write_state_pages_to_cpu( - self, - page_kind: PageKind, - page_slot_ids: List[int], - cpu_page_buffers: torch.Tensor, - cpu_page_indices: List[int], - layout: StateTransferLayout, - cpu_page_ready_flags: list[bool] | None = None, - ) -> None: - """将指定页池槽位的快照写入 CPU 页。""" - assert layout.state_name == self.state_name, "transfer layout does not describe this resource" - assert 0 <= layout.byte_start < layout.byte_end, "invalid State byte interval" - assert cpu_page_buffers.device.type == "cpu", "CPU cache buffer must reside on CPU" - pool = self._pool(page_kind) - assert isinstance(page_slot_ids, list), "slot indices must be a list" - assert isinstance(cpu_page_indices, list), "CPU page indices must be a list" - assert len(page_slot_ids) == len(cpu_page_indices), "CPU page and slot counts differ" - if cpu_page_ready_flags is not None: - assert len(cpu_page_ready_flags) == len(cpu_page_indices), "CPU page and ready counts differ" - page_fields, page_indices = [], [] - for index, (slot, page) in enumerate(zip(page_slot_ids, cpu_page_indices)): - # ready 页不参与复制；同一传输区域的所有页面组成一个同形状批次。 - if cpu_page_ready_flags is not None and cpu_page_ready_flags[index]: - continue - page_fields.append(self._fields(pool.conv_state[slot], pool.ssm_state[slot])) - page_indices.append(page) - for copies in bounded_att_buffer_copies( - page_fields=page_fields, - page_buffers=cpu_page_buffers, - layout=layout, - page_indices=page_indices, - write=True, - ): - copy_bounded_att_tensors(copies=copies) - - def read_state_pages_from_cpu( - self, - cpu_page_buffers: torch.Tensor, - cpu_page_indices: List[int], - big_page_slot_ids: List[int], - layout: StateTransferLayout, - ) -> None: - """读入指定大页槽位；请求运行态另行通过 restore_bounded_att_runtime_from_pages 恢复。""" - assert layout.state_name == self.state_name, "transfer layout does not describe this resource" - assert 0 <= layout.byte_start < layout.byte_end, "invalid State byte interval" - assert cpu_page_buffers.device.type == "cpu", "CPU cache buffer must reside on CPU" - pool = self.big_pages_pool - assert isinstance(big_page_slot_ids, list), "slot indices must be a list" - assert isinstance(cpu_page_indices, list), "CPU page indices must be a list" - assert len(cpu_page_indices) == len(big_page_slot_ids), "CPU page and slot counts differ" - # 本地目标：选定大页槽位的 Conv/SSM 快照，保留本 TP rank 的分片形状。 - pool_page_fields = [self._fields(pool.conv_state[slot], pool.ssm_state[slot]) for slot in big_page_slot_ids] - # 外部来源：layout.regions 描述 CPU 全局字段和 TP 范围；write=False 将其读入本地页池。 - # 窗口保持快照顺序；请求 ring 的位置由后续 restore_bounded_att_runtime_from_pages 按长度恢复。 - for copies in bounded_att_buffer_copies( - page_fields=pool_page_fields, - page_buffers=cpu_page_buffers, - layout=layout, - page_indices=cpu_page_indices, - write=False, - ): - copy_bounded_att_tensors(copies=copies) - - # ======================================================================== - # 四、DP 迁移：大页快照与基准运行态 - # ======================================================================== - - def copy_big_pages_from_dp( - self, - src_page_ids: List[int], - dst_page_ids: List[int], - remote_dp_bounded_state: KDAState, - ) -> None: - """复制同一 TP 分片的大页 Conv/SSM 快照，保留源页和双方请求运行态。""" - assert isinstance(src_page_ids, list) and isinstance(dst_page_ids, list), "page indices must be lists" - assert len(src_page_ids) == len(dst_page_ids), "source and destination page counts differ" - assert ( - type(remote_dp_bounded_state) is type(self) and remote_dp_bounded_state.config == self.config - ), "DP copy requires the same algorithm and equivalent TP state configurations" - src_pool = remote_dp_bounded_state.big_pages_pool - dst_pool = self.big_pages_pool - assert src_pool is not None and dst_pool is not None, "DP big page copy requires both page pools" - # 大页保存的是基准快照，与请求槽位及 MTP 候选区无关；页的引用和使用期限由调用方管理。 - # 页池均须在本进程注册；GPU 在当前流直接访问两端 CPU 存储，完成后才能归还页面。 - conv_copies, ssm_copies = [], [] - for src_page, dst_page in zip(src_page_ids, dst_page_ids): - conv_copies.append((src_pool.conv_state[src_page], dst_pool.conv_state[dst_page])) - ssm_copies.append((src_pool.ssm_state[src_page], dst_pool.ssm_state[dst_page])) - copy_bounded_att_tensors(copies=conv_copies) - copy_bounded_att_tensors(copies=ssm_copies) - - def copy_runtime_state_from_dp( - self, - src_req_indices: List[int], - dst_req_indices: List[int], - remote_dp_bounded_state: KDAState, - ) -> None: - """复制同一 TP 分片的 prefill 基准运行态并重置目标游标，不写入页池。""" - assert isinstance(src_req_indices, list) and isinstance(dst_req_indices, list), "request indices must be lists" - assert len(src_req_indices) == len(dst_req_indices), "source and destination request counts differ" - assert ( - type(remote_dp_bounded_state) is type(self) and remote_dp_bounded_state.config == self.config - ), "DP copy requires the same algorithm and equivalent TP state configurations" - conv_copies, ssm_copies = [], [] - for src_req, dst_req in zip(src_req_indices, dst_req_indices): - source = remote_dp_bounded_state.get_accepted_state(src_req) - target = self.prepare_restore(dst_req) - conv_copies.append((source.conv, target.conv)) - ssm_copies.append((source.ssm, target.ssm)) - copy_bounded_att_tensors(copies=conv_copies) - copy_bounded_att_tensors(copies=ssm_copies) - - # ======================================================================== - # 五、PD 分离传输 - # ======================================================================== - - @deepcopy_result - @lru_cache(maxsize=None) - def get_pd_transfer_regions(self) -> tuple[TransferRegion, ...]: - """返回 PD 传输字段与切片，交由管理器安排其在联合页中的位置。""" - return self._get_transfer_regions() - - def write_to_pd_page( - self, - req_idx: int, - pd_page_buffers: torch.Tensor, - pd_page_index: int, - state_token_len: int, - layout: StateTransferLayout, - ) -> None: - """写出 state_token_len 处的已接受状态，不修改源 Replay 历史。""" - assert layout.state_name == self.state_name, "transfer layout does not describe this resource" - assert 0 <= layout.byte_start < layout.byte_end, "invalid State byte interval" - fields = self._fields(*self.get_accepted_state(req_idx)) - for copies in bounded_att_buffer_copies( - page_fields=[fields], - page_buffers=pd_page_buffers, - layout=layout, - page_indices=[pd_page_index], - write=True, - ): - copy_bounded_att_tensors(copies=copies) - - def read_from_pd_page( - self, - pd_page_buffers: torch.Tensor, - req_idx: int, - pd_page_index: int, - state_token_len: int, - layout: StateTransferLayout, - ) -> None: - """恢复 state_token_len 处的基准状态并规范化控制信息；Conv/SSM 不依赖长度定位。""" - assert layout.state_name == self.state_name, "transfer layout does not describe this resource" - assert 0 <= layout.byte_start < layout.byte_end, "invalid State byte interval" - fields = self._fields(*self.prepare_restore(req_idx)) - for copies in bounded_att_buffer_copies( - page_fields=[fields], - page_buffers=pd_page_buffers, - layout=layout, - page_indices=[pd_page_index], - write=False, - ): - copy_bounded_att_tensors(copies=copies) - if self.mtp_state_index is not None: - self.mtp_state_index[req_idx].fill_(0) - - # ======================================================================== - # 六、私有辅助实现 - # ======================================================================== - - def _get_transfer_regions(self) -> tuple[TransferRegion, ...]: - """依次描述全局 Conv 和 SSM 区域，并标明当前 TP rank 的读写范围。""" - # 阶段一：确定全局 Conv 形状；只传有效历史窗口，不包含 MTP 候选区和游标。 - config = self.config - storage_layer_count = len(config.logical_layer_ids) - conv_channels = 2 * config.global_k_heads * config.key_dim + config.global_v_heads * config.value_dim - conv_shape = (storage_layer_count, conv_channels, config.conv_width) - conv_byte_end = prod(conv_shape) * config.conv_dtype.itemsize - regions = [] - logger.debug( - "%s building transfer regions: TP=%d/%d, Conv shape=%s", - self.state_name, - config.tp_rank, - config.tp_world_size, - conv_shape, - ) - - # 阶段二：Conv 按 Q/K/V 拼接；三个映射共用同一块全局 Conv 区域，不能重复累计字节数。 - local_channel_start = 0 - global_channel_start = 0 - for name, local_head_count, global_head_count, head_dim in ( - ("Q", config.num_k_heads, config.global_k_heads, config.key_dim), - ("K", config.num_k_heads, config.global_k_heads, config.key_dim), - ("V", config.num_v_heads, config.global_v_heads, config.value_dim), - ): - # 配置已保证完整分片或副本。每组连续 rank 保存相同 heads，仅组内首 rank 写出。 - replica_count = config.tp_world_size * local_head_count // global_head_count - head_start = config.tp_rank // replica_count * local_head_count - write_enabled = config.tp_rank % replica_count == 0 - local_channel_count = local_head_count * head_dim - target_channel_start = global_channel_start + head_start * head_dim - conv_region = TransferRegion( - field="conv", - byte_start=0, - byte_end=conv_byte_end, - dtype=config.conv_dtype, - axes=("layer", "channel", "window"), - shape=conv_shape, - axis_mappings=( - AxisMapping( - axis="channel", - local_start=local_channel_start, - external_start=target_channel_start, - length=local_channel_count, - ), - AxisMapping( - axis="window", - local_start=0, - external_start=0, - length=config.conv_width, - ), - ), - write_enabled=write_enabled, - ) - regions.append(conv_region) - logger.debug( - "%s Conv %s: local_channel_start=%d, global_channel_start=%d, " - "channels=%d, replicas=%d, write_enabled=%s", - self.state_name, - name, - local_channel_start, - target_channel_start, - local_channel_count, - replica_count, - write_enabled, - ) - local_channel_start += local_channel_count - global_channel_start += global_head_count * head_dim - # 阶段三：SSM 沿 V head 切分；接在 Conv 之后，按 SSM dtype 对齐字节起点。 - replica_count = config.tp_world_size * config.num_v_heads // config.global_v_heads - head_start = config.tp_rank // replica_count * config.num_v_heads - write_enabled = config.tp_rank % replica_count == 0 - ssm_shape = (storage_layer_count, config.global_v_heads, config.key_dim, config.value_dim) - element_bytes = config.ssm_dtype.itemsize - ssm_byte_start = (conv_byte_end + element_bytes - 1) // element_bytes * element_bytes - ssm_region = TransferRegion( - field="ssm", - byte_start=ssm_byte_start, - byte_end=ssm_byte_start + prod(ssm_shape) * element_bytes, - dtype=config.ssm_dtype, - axes=("layer", "head", "key_dim", "value_dim"), - shape=ssm_shape, - axis_mappings=( - AxisMapping( - axis="head", - local_start=0, - external_start=head_start, - length=config.num_v_heads, - ), - AxisMapping( - axis="key_dim", - local_start=0, - external_start=0, - length=config.key_dim, - ), - AxisMapping( - axis="value_dim", - local_start=0, - external_start=0, - length=config.value_dim, - ), - ), - write_enabled=write_enabled, - ) - regions.append(ssm_region) - logger.debug( - "%s transfer regions built: SSM head_start=%d, heads=%d, write_enabled=%s, SSM bytes=[%d, %d)", - self.state_name, - head_start, - config.num_v_heads, - write_enabled, - ssm_region.byte_start, - ssm_region.byte_end, - ) - return tuple(regions) - def _pool(self, page_kind: PageKind): - """根据大小页类型取得已启用的页池，不访问 allocator。""" - assert page_kind in ("big", "small"), f"unknown page kind {page_kind}" - pool = self.big_pages_pool if page_kind == "big" else self.small_pages_pool - assert pool is not None, f"{page_kind} pages are disabled for {self.state_name}" - return pool +class KDAState(ConvSSMState): - @staticmethod - def _fields(conv: torch.Tensor, ssm: torch.Tensor): - return { - "conv": TensorField(tensor=conv, axes=("layer", "channel", "window")), - "ssm": TensorField(tensor=ssm, axes=("layer", "head", "key_dim", "value_dim")), - } + algorithm = "kda" + pages_pool_class = KDAPagesPool diff --git a/lightllm/common/attention_state/bounded_att/sliding_window/fp8_e8m0_swa_state.py b/lightllm/common/attention_state/bounded_att/sliding_window/fp8_e8m0_swa_state.py index a1d72102ac..5447ac3740 100644 --- a/lightllm/common/attention_state/bounded_att/sliding_window/fp8_e8m0_swa_state.py +++ b/lightllm/common/attention_state/bounded_att/sliding_window/fp8_e8m0_swa_state.py @@ -179,7 +179,9 @@ def restore_bounded_att_runtime_from_pages( write=False, ) - def commit_bounded_att_state(self, req_indices: torch.Tensor, accepted_steps: torch.Tensor) -> None: + def commit_bounded_att_state( + self, req_indices: torch.Tensor, accepted_steps: torch.Tensor, *, context=None + ) -> None: """ring 按绝对 token 位置寻址，无需提交额外游标。""" pass diff --git a/lightllm/common/attention_state/bounded_att/ssm/__init__.py b/lightllm/common/attention_state/bounded_att/ssm/__init__.py new file mode 100644 index 0000000000..a2d69e470c --- /dev/null +++ b/lightllm/common/attention_state/bounded_att/ssm/__init__.py @@ -0,0 +1,34 @@ +import torch + + +def create_ssm_state(config, request_capacity, mtp_step, device, mode, replay_config, *, kda, gate_lower_bound): + if mode not in ("native", "compact", "replay"): + raise ValueError(f"unknown SSM state mode: {mode}") + slots = mtp_step + 1 if mode == "native" else 1 + state = torch.zeros( + (len(config.logical_layer_ids), request_capacity * slots, config.num_v_heads, config.key_dim, config.value_dim), + dtype=config.ssm_dtype, + device=device, + ) + kwargs = dict(kda=kda, gate_lower_bound=gate_lower_bound) + if mode == "native": + from .native import NativeSSMState + + return NativeSSMState(state, mtp_step + 1, **kwargs) + kwargs.update(num_key_heads=config.num_k_heads) + if mode == "compact": + from .compact import CompactSSMState + + return CompactSSMState(state, mtp_step + 1, config.conv_dtype, **kwargs) + from .replay import ReplaySSMState + + replay_config = replay_config or {} + return ReplaySSMState( + state, + replay_config.get("cache_len", 16), + mtp_step + 1, + config.conv_dtype, + compute_mode=replay_config.get("compute_mode", "inline"), + fold_method=replay_config.get("fold_method", "raw"), + **kwargs, + ) diff --git a/lightllm/common/attention_state/bounded_att/ssm/base.py b/lightllm/common/attention_state/bounded_att/ssm/base.py new file mode 100644 index 0000000000..ceb719edf0 --- /dev/null +++ b/lightllm/common/attention_state/bounded_att/ssm/base.py @@ -0,0 +1,66 @@ +from lightllm.common.basemodel.triton_kernel.mtp_utils import linear_att_mtp_state_index_update + + +def update_accepted_indices(reqs, accepted_steps, indices, context, *, replay_cursors=None): + if context is None: + indices[reqs] = accepted_steps - 1 + else: + linear_att_mtp_state_index_update( + req_to_mtp_state_index=indices, + b_req_mtp_start_loc=context.b_req_mtp_start_loc, + b_req_idx=context.b_req_idx, + b_mtp_index=context.b_mtp_index, + accepted_index=context.accepted_index, + verify_width=context.verify_width, + replay_cursors=replay_cursors, + ) + + +class SSMState: + slots_per_request = 1 + + def clear_history(self, req): + pass + + def prepare_decode(self, reqs, cu_seqlens=None): + pass + + def accept_single_token(self, reqs): + pass + + def merge_accepted_updates(self, reqs): + pass + + def decode_indices(self, reqs, speculative): + return reqs + + def get_accepted_state(self, req_idx, accepted_index=None): + return self.state[:, req_idx] + + def create_commit(self, run_config): + return None + + def validate_commit(self, context, state_name): + pass + + def commit(self, reqs, accepted_steps, indices, context, state_name): + update_accepted_indices(reqs, accepted_steps, indices, context) + + def forward( + self, + layer, + q, + k, + v, + a, + b, + a_log, + bias, + reqs, + cu_seqlens=None, + run_config=None, + *, + state_indices=None, + accepted_token_counts=None, + ): + raise NotImplementedError diff --git a/lightllm/common/attention_state/bounded_att/ssm/compact.py b/lightllm/common/attention_state/bounded_att/ssm/compact.py new file mode 100644 index 0000000000..7d570f93cd --- /dev/null +++ b/lightllm/common/attention_state/bounded_att/ssm/compact.py @@ -0,0 +1,177 @@ +import torch +import triton + +from lightllm.common.ssm_types import SSMRunConfig +from lightllm.common.basemodel.triton_kernel.linear_att.compact_ssm import _compact +from lightllm.common.basemodel.triton_kernel.linear_att.ssm_autotune import resolve_run_config +from .base import SSMState, update_accepted_indices + + +class CompactSSMState(SSMState): + def __init__( + self, + state, + verify_width, + activation_dtype, + *, + num_key_heads=None, + kda=False, + gate_lower_bound=-5.0, + run_config=None, + ): + self.state = state + self.kda = kda + self.gate_lower_bound = gate_lower_bound + self.verify_width = verify_width + # Verify and accepted-state reconstruction must share the same layout. + self._config_is_fixed = run_config is not None + self.batch_configs = {} + self.run_config = dict(run_config or {"BV": 8, "num_warps": 1}) + assert self.run_config["BV"] in (8, 16, 32, 64, 128) + assert self.run_config["num_warps"] in (1, 2, 4, 8) + layers, slots, hv, k, v = state.shape + self.num_key_heads = hv if num_key_heads is None else num_key_heads + assert self.num_key_heads > 0 and hv % self.num_key_heads == 0 + self.hold_req_id = slots - 1 + shape = (layers, slots, hv, verify_width) + self.keys = torch.empty( + (layers, slots, self.num_key_heads, verify_width, 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 kda else shape, device=state.device, dtype=torch.float32) + self.betas = torch.empty(shape, device=state.device, dtype=torch.float32) + + def get_run_config(self, q, k, v, a, b, a_log, bias, cu_seqlens=None) -> SSMRunConfig: + config = resolve_run_config(self, "compact", q, k, v, a, b, a_log, bias, cu_seqlens) + return SSMRunConfig(**config) + + def forward( + self, + linear_layer_idx, + q, + k, + v, + a, + b, + a_log, + bias, + reqs, + cu_seqlens=None, + run_config=None, + *, + state_indices=None, + accepted_token_counts=None, + ): + assert cu_seqlens is not None, "compact replay is only used for speculative verify" + config = run_config or self.get_run_config(q, k, v, a, b, a_log, bias, cu_seqlens) + if not isinstance(config, SSMRunConfig): + config = SSMRunConfig(**config) + _, slots, hv, kd, vd = self.state.shape + assert self.num_key_heads in (q.shape[-2], hv) + out = torch.empty_like(v) + # Match native reduction layouts: changing BV can alter BF16 rounding. + bv = config.BV + _compact[(triton.cdiv(vd, bv), reqs.numel(), hv)]( + q, + k, + v, + a, + b, + a_log, + bias, + self.state[linear_layer_idx], + self.keys[linear_layer_idx], + self.values[linear_layer_idx], + self.decays[linear_layer_idx], + self.betas[linear_layer_idx], + 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, + self.num_key_heads, + kd, + vd, + self.verify_width, + self.hold_req_id, + triton.next_power_of_2(kd), + bv, + False, + self.kda, + self.gate_lower_bound, + num_warps=config.num_warps, + num_stages=config.num_stages, + ) + return out, config + + def accept_updates(self, reqs, accepted_state_indices, run_config=None): + """Merge this round's prefix; accepted_state_indices contains per-request last accepted indexes.""" + assert run_config is not None or not self.batch_configs, "compact acceptance requires the verify configuration" + layers, slots, hv, kd, vd = self.state.shape + config = run_config or self.run_config + if not isinstance(config, SSMRunConfig): + config = SSMRunConfig(**config) + bv = config.BV + _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_state_indices, + None, + 0, + 0, + 0, + 0, + 0, + slots, + hv, + hv, + self.num_key_heads, + kd, + vd, + self.verify_width, + self.hold_req_id, + triton.next_power_of_2(kd), + bv, + True, + self.kda, + self.gate_lower_bound, + num_warps=config.num_warps, + num_stages=config.num_stages, + ) + + def create_commit(self, run_config): + if not isinstance(run_config, SSMRunConfig): + raise ValueError("compact commit requires the configuration of its forward") + return run_config + + def validate_commit(self, context, state_name): + if context is None or not context.batches: + raise ValueError("compact acceptance requires forward commit batches") + for _, config in context.partitions(state_name): + if not isinstance(config, SSMRunConfig): + raise ValueError(f"missing compact forward configuration for {state_name}") + + def commit(self, reqs, accepted_steps, indices, context, state_name): + self.validate_commit(context, state_name) + update_accepted_indices(reqs, accepted_steps, indices, context) + for partition, config in context.partitions(state_name): + self.accept_updates(reqs[partition], indices, run_config=config) diff --git a/lightllm/common/attention_state/bounded_att/ssm/forward.py b/lightllm/common/attention_state/bounded_att/ssm/forward.py new file mode 100644 index 0000000000..b827149a36 --- /dev/null +++ b/lightllm/common/attention_state/bounded_att/ssm/forward.py @@ -0,0 +1,64 @@ +from dataclasses import dataclass +from typing import Optional, TYPE_CHECKING + +import torch + +from lightllm.common.attention_state.commit import StateCommit +from lightllm.common.ssm_types import SSMRunConfig + +if TYPE_CHECKING: + from ..conv_ssm_state import ConvSSMState + + +@dataclass +class SSMForward: + resource: "ConvSSMState" + req_indices: torch.Tensor + state_indices: torch.Tensor + cu_seqlens: Optional[torch.Tensor] + accepted_token_counts: Optional[torch.Tensor] + run_config: Optional[SSMRunConfig] = None + + def run_layer(self, logical_layer, q, k, v, a, b, a_log, bias): + output, self.run_config = self.resource._ssm.forward( + self.resource.get_storage_layer_index(logical_layer), + q, + k, + v, + a, + b, + a_log, + bias, + self.req_indices, + self.cu_seqlens, + run_config=self.run_config, + state_indices=self.state_indices, + accepted_token_counts=self.accepted_token_counts, + ) + return output + + def finish(self): + payload = self.resource._ssm.create_commit(self.run_config) + if payload is None: + return None + if self.resource.state_name is None: + raise ValueError("forward commits require a registered resource") + return StateCommit(self.resource.state_name, payload) + + def copy_inputs(self, source: "SSMForward"): + if self.resource is not source.resource: + raise ValueError("cannot copy a forward for a different resource") + copies = [] + for name in ("req_indices", "state_indices", "cu_seqlens", "accepted_token_counts"): + target, value = getattr(self, name), getattr(source, name) + if (target is None) != (value is None): + raise ValueError(f"captured SSM layout changed: {name}") + if target is not None: + if target.shape != value.shape or target.dtype != value.dtype or target.device != value.device: + raise ValueError(f"captured SSM tensor layout changed: {name}") + copies.append((target, value)) + copied = set() + for target, value in copies: + if target is not None and target.data_ptr() != value.data_ptr() and target.data_ptr() not in copied: + target.copy_(value, non_blocking=True) + copied.add(target.data_ptr()) diff --git a/lightllm/common/attention_state/bounded_att/ssm/native.py b/lightllm/common/attention_state/bounded_att/ssm/native.py new file mode 100644 index 0000000000..b090a3131d --- /dev/null +++ b/lightllm/common/attention_state/bounded_att/ssm/native.py @@ -0,0 +1,96 @@ +import torch + +from .base import SSMState + + +class NativeSSMState(SSMState): + def __init__(self, state, verify_width, *, kda=False, gate_lower_bound=-5.0): + self.state = state + self.slots_per_request = verify_width + self.kda = kda + self.gate_lower_bound = gate_lower_bound + + def decode_indices(self, reqs, speculative): + if not speculative: + return reqs if self.slots_per_request == 1 else reqs * self.slots_per_request + offsets = torch.arange(self.slots_per_request, device=reqs.device, dtype=reqs.dtype) + return reqs[:, None] * self.slots_per_request + offsets[None, :] + + def get_accepted_state(self, req_idx, accepted_index=None): + if self.slots_per_request == 1: + return self.state[:, req_idx] + index = req_idx * self.slots_per_request + accepted_index + return self.state.index_select(1, index.long().view(1))[:, 0] + + def forward( + self, + layer, + q, + k, + v, + a, + b, + a_log, + bias, + reqs, + cu_seqlens=None, + run_config=None, + *, + state_indices=None, + accepted_token_counts=None, + ): + state = self.state[layer] + if self.kda: + from lightllm.common.basemodel.triton_kernel.linear_att.fla.ops.kda_decode import fused_recurrent_kda + + output, _ = fused_recurrent_kda( + q=q, + k=k, + v=v, + raw_gate=a.view(*q.shape[:2], -1), + raw_beta=b.view(*q.shape[:2], -1), + a_log=a_log, + gate_bias=bias, + initial_state=state, + lower_bound=self.gate_lower_bound, + inplace_final_state=True, + ssm_state_indices=state_indices, + cu_seqlens=cu_seqlens, + num_accepted_tokens=accepted_token_counts, + ) + elif cu_seqlens is not None: + from lightllm.common.basemodel.triton_kernel.linear_att.mtp_fused_recurrent import ( + mtp_fused_recurrent_gated_delta_rule, + ) + + output, _ = mtp_fused_recurrent_gated_delta_rule( + q=q, + k=k, + v=v, + initial_state=state, + cu_seqlens=cu_seqlens.to(torch.long), + ssm_state_indices=state_indices, + ssm_state_write_indices=state_indices, + num_accepted_tokens=accepted_token_counts, + A_log=a_log, + dt_bias=bias, + a_raw=a, + b_raw=b, + ) + else: + from lightllm.common.basemodel.triton_kernel.linear_att.fla.ops import fused_recurrent_gated_delta_rule + + output, _ = fused_recurrent_gated_delta_rule( + q=q, + k=k, + v=v, + initial_state=state, + inplace_final_state=True, + ssm_state_indices=state_indices, + use_qk_l2norm_in_kernel=True, + A_log=a_log, + dt_bias=bias, + a_raw=a, + b_raw=b, + ) + return output, None diff --git a/lightllm/common/attention_state/bounded_att/ssm/replay.py b/lightllm/common/attention_state/bounded_att/ssm/replay.py new file mode 100644 index 0000000000..ce3f62f69f --- /dev/null +++ b/lightllm/common/attention_state/bounded_att/ssm/replay.py @@ -0,0 +1,300 @@ +import torch +import triton + +from lightllm.common.ssm_types import SSMRunConfig +from lightllm.common.basemodel.triton_kernel.linear_att.replayssm import ( + _accept_updates, + _prepare_fold, + _fold_history, + _materialize, + _replay, +) +from lightllm.common.basemodel.triton_kernel.linear_att.ssm_autotune import resolve_run_config +from .base import SSMState, update_accepted_indices + + +class ReplaySSMState(SSMState): + """Request-owned scratch; canonical CPU/PD checkpoints never include it.""" + + def __init__( + self, + state, + capacity, + verify_width, + activation_dtype=torch.bfloat16, + *, + num_key_heads=None, + kda=False, + gate_lower_bound=-5.0, + compute_mode="inline", + run_config=None, + fold_method="raw", + ): + assert state.dtype in (torch.float32, torch.bfloat16), "ReplaySSM requires FP32 or BF16 SSM state" + assert verify_width > 0 and capacity >= max(4, verify_width) and capacity & (capacity - 1) == 0 + assert compute_mode in ("inline", "precompute", "parallel") + assert fold_method in ("raw", "derived") + assert fold_method == "raw" or compute_mode == "parallel" + if compute_mode == "parallel": + assert not kda and verify_width > 1, "Parallel ReplaySSM requires GDN MTP" + assert state.shape[-2:] == (128, 128), "Parallel ReplaySSM requires 128x128 heads" + assert activation_dtype == torch.bfloat16, "Parallel ReplaySSM requires BF16 activations" + assert state.dtype == torch.bfloat16, "Parallel ReplaySSM currently supports BF16 state" + self.fold_method = fold_method + self.state = state + self.kda = kda + self.gate_lower_bound = gate_lower_bound + self.fold_programs = ( + 8 * torch.cuda.get_device_properties(state.device).multi_processor_count if state.is_cuda else 0 + ) + self.capacity = capacity + self.verify_width = verify_width + self.compute_mode = compute_mode + self.run_config = run_config + self._config_is_fixed = run_config is not None + self.batch_configs = {} + layers, slots, hv, k, v = state.shape + # Omission preserves the redundant layout for existing callers and A/B. + self.num_key_heads = hv if num_key_heads is None else num_key_heads + assert self.num_key_heads > 0 and hv % self.num_key_heads == 0 + self.hold_req_id = slots - 1 + self.cursors = torch.zeros(slots, dtype=torch.int32, device=state.device) + self.forward_cursors = torch.zeros_like(self.cursors) + # Alternate halves on fold so CTAs cannot overwrite history another + # V tile is still reading. Cursor packs phase (2*L) and count (0..L). + # Derived corrections can exceed the activation/state range even when + # the recurrence stays finite. Only BF16 has enough exponent range for + # the short-history narrowing optimization; FP16 must retain FP32. + history_dtype = ( + activation_dtype + if activation_dtype == torch.bfloat16 + and verify_width > 1 + and compute_mode == "precompute" + and capacity <= 8 + else torch.float32 + ) + key_shape = (layers, slots, self.num_key_heads, 2 * capacity, k) + self.keys = torch.empty(key_shape, dtype=history_dtype, device=state.device) + self.deltas = torch.empty((layers, slots, hv, 2 * capacity, v), dtype=history_dtype, device=state.device) + self.gates = torch.empty( + (layers, slots, hv, 2 * capacity, k) if kda else (layers, slots, hv, 2 * capacity), + dtype=torch.float32, + device=state.device, + ) + self.raw_keys = torch.empty(key_shape, 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 clear_history(self, req): + """Invalidate history for a reused/restored request; leave its SSM state intact.""" + if type(req) is int: + # Scalar assignment copies a CPU tensor into the CUDA slot. + self.cursors[req].zero_() + else: + self.cursors[req] = 0 + + def prepare_decode(self, reqs, cu_seqlens=None): + """Snapshot history for all layers and prepare the base for accepted updates.""" + positions = torch.empty_like(reqs) + layers, _, hv, kd, vd = self.state.shape + active = torch.empty((reqs.numel() + 1,), dtype=torch.int32, device=reqs.device) + _prepare_fold[(1,)]( + reqs, + self.cursors, + positions, + self.forward_cursors, + cu_seqlens, + active, + reqs.numel(), + self.hold_req_id, + self.capacity, + cu_seqlens is not None, + triton.next_power_of_2(reqs.numel()), + self.verify_width, + self.compute_mode == "parallel", + num_warps=4, + ) + programs = min(self.fold_programs, triton.cdiv(vd, 32) * reqs.numel() * layers * hv) + _fold_history[(programs,)]( + self.state, + self.raw_keys, + self.raw_values, + self.gates, + self.betas, + reqs, + positions, + active, + reqs.numel(), + layers, + hv, + self.num_key_heads, + kd, + vd, + self.capacity, + triton.next_power_of_2(kd), + 32, + self.state.stride(0), + self.raw_keys.stride(0), + self.raw_values.stride(0), + self.gates.stride(0), + self.betas.stride(0), + self.kda, + self.keys, + self.deltas, + self.fold_method == "derived", + BL=max(16, self.capacity), + num_warps=2, + num_stages=1, + ) + return positions + + def accept_updates(self, reqs, accepted_state_indices=None): + """Accept this round's prefix; accepted_state_indices contains per-request last accepted indexes.""" + _accept_updates[(triton.cdiv(reqs.numel(), 256),)]( + self.cursors, + reqs, + accepted_state_indices, + reqs.numel(), + self.hold_req_id, + self.verify_width, + 256, + ) + + def merge_accepted_updates(self, reqs): + """Fold accepted history into SSM state and clear its history positions.""" + if reqs.numel() == 0: + return + self._materialize_accepted_state(reqs, self.state, snapshot=False) + self.cursors.index_fill_(0, reqs.long(), 0) + + def _materialize_accepted_state(self, reqs, output, snapshot): + layers, slots, hv, k, v = self.state.shape + config = self.run_config or {"BV": 32, "num_warps": 4} + bv = config["BV"] + _materialize[(triton.cdiv(v, bv), reqs.numel(), layers * hv)]( + self.state, + output, + self.raw_keys, + self.raw_values, + self.gates, + self.betas, + self.cursors, + reqs, + slots, + hv, + self.num_key_heads, + k, + v, + self.capacity, + self.hold_req_id, + triton.next_power_of_2(k), + bv, + snapshot, + self.kda, + num_warps=config["num_warps"], + num_stages=config.get("num_stages", 3), + ) + + def get_accepted_state(self, req_idx, accepted_index=None): + """Return a canonical checkpoint without changing the active state or history.""" + layers, _, hv, k, v = self.state.shape + output = torch.empty((layers, hv, k, v), dtype=self.state.dtype, device=self.state.device) + reqs = torch.tensor([req_idx], dtype=torch.int32, device=self.state.device) + self._materialize_accepted_state(reqs, output, snapshot=True) + return output + + def get_run_config(self, q, k, v, a, b, a_log, bias, cu_seqlens=None) -> SSMRunConfig: + if self.compute_mode == "parallel": + config = self.run_config or {"BV": 128, "num_warps": 2, "num_stages": 1} + else: + config = resolve_run_config(self, "replay", q, k, v, a, b, a_log, bias, cu_seqlens) + return SSMRunConfig(**config) + + def forward( + self, + linear_layer_idx, + q, + k, + v, + a, + b, + a_log, + bias, + reqs, + cu_seqlens=None, + run_config=None, + *, + state_indices=None, + accepted_token_counts=None, + ): + config = run_config or self.get_run_config(q, k, v, a, b, a_log, bias, cu_seqlens) + if not isinstance(config, SSMRunConfig): + config = SSMRunConfig(**config) + hv, kd, vd = self.state.shape[-3:] + assert self.num_key_heads in (q.shape[-2], hv) + axis = 1 if cu_seqlens is not None else 0 + out = torch.empty_like(v) + bv = config.BV + precompute_projections = self.verify_width > 1 and bv >= 16 and self.compute_mode in ("precompute", "parallel") + if self.compute_mode == "parallel": + assert precompute_projections, "Parallel ReplaySSM requires BV >= 16" + # Resolve tile dimensions in Python: bare max() in Triton JIT code + # lowers to a tensor operation, which cannot define static shapes. + bw = triton.next_power_of_2(self.verify_width) + _replay[(triton.cdiv(vd, bv), reqs.numel(), hv)]( + q, + k, + v, + a, + b, + a_log, + bias, + self.state[linear_layer_idx], + self.keys[linear_layer_idx], + self.deltas[linear_layer_idx], + self.gates[linear_layer_idx], + self.raw_keys[linear_layer_idx], + self.raw_values[linear_layer_idx], + self.betas[linear_layer_idx], + reqs, + self.forward_cursors, + cu_seqlens, + out, + q.stride(axis), + k.stride(axis), + v.stride(axis), + a.stride(0), + b.stride(0), + q.shape[-2], + hv, + self.num_key_heads, + kd, + vd, + self.capacity, + self.hold_req_id, + triton.next_power_of_2(kd), + bv, + cu_seqlens is not None, + precompute_projections, + bw, + reqs.numel() > 12 and self.compute_mode != "parallel", + self.kda, + self.gate_lower_bound, + self.compute_mode == "parallel", + BQK=max(16, 2 * bw), + BR=triton.next_power_of_2(max(16, self.capacity + bw)), + num_warps=config.num_warps, + num_stages=config.num_stages, + ) + return out, config + + def accept_single_token(self, reqs): + + self.accept_updates(reqs) + + def commit(self, reqs, accepted_steps, indices, context, state_name): + if context is not None and self.verify_width > 1: + update_accepted_indices(reqs, accepted_steps, indices, context, replay_cursors=self.cursors) + else: + update_accepted_indices(reqs, accepted_steps, indices, context) + self.accept_updates(reqs, indices) diff --git a/lightllm/common/attention_state/bounded_att/window_compressor/window_compressor_state.py b/lightllm/common/attention_state/bounded_att/window_compressor/window_compressor_state.py index cea0a49af2..2b937bf548 100644 --- a/lightllm/common/attention_state/bounded_att/window_compressor/window_compressor_state.py +++ b/lightllm/common/attention_state/bounded_att/window_compressor/window_compressor_state.py @@ -210,7 +210,9 @@ def restore_bounded_att_runtime_from_pages( write=False, ) - def commit_bounded_att_state(self, req_indices: torch.Tensor, accepted_steps: torch.Tensor) -> None: + def commit_bounded_att_state( + self, req_indices: torch.Tensor, accepted_steps: torch.Tensor, *, context=None + ) -> None: """按绝对位置寻址；接受后继续覆盖 ring，不维护额外 MTP 游标。""" pass diff --git a/lightllm/common/attention_state/commit.py b/lightllm/common/attention_state/commit.py new file mode 100644 index 0000000000..778e154eb5 --- /dev/null +++ b/lightllm/common/attention_state/commit.py @@ -0,0 +1,59 @@ +from dataclasses import dataclass +from typing import Tuple + +import torch + + +@dataclass(frozen=True) +class StateCommit: + state_name: str + payload: object + + +@dataclass(frozen=True) +class StateCommitBatch: + + request_count: int + commits: Tuple[StateCommit, ...] = () + + def __post_init__(self): + if type(self.request_count) is not int or self.request_count < 0: + raise ValueError("commit batch requires a nonnegative request count") + names = [commit.state_name for commit in self.commits] + if len(set(names)) != len(names): + raise ValueError("duplicate state commit in a forward") + + +@dataclass(frozen=True) +class MtpCommitContext: + + b_req_mtp_start_loc: torch.Tensor + b_req_idx: torch.Tensor + b_mtp_index: torch.Tensor + accepted_index: torch.Tensor + verify_width: int + batches: Tuple[StateCommitBatch, ...] = () + + def validate(self, request_count: int, state_names=None): + if type(self.verify_width) is not int or self.verify_width <= 0: + raise ValueError("verification width must be positive") + rows = (self.b_req_idx, self.b_mtp_index, self.accepted_index) + if any(row.ndim != 1 or row.numel() != self.b_req_idx.numel() for row in rows): + raise ValueError("verification row tensors must have the same one-dimensional layout") + if self.b_req_mtp_start_loc.numel() != request_count: + raise ValueError("verification and logical request counts differ") + if self.batches and sum(batch.request_count for batch in self.batches) != request_count: + raise ValueError("commit batches must partition all unpadded requests") + for batch in self.batches: + for commit in batch.commits: + if state_names is not None and commit.state_name not in state_names: + raise ValueError(f"commit belongs to an unknown state: {commit.state_name}") + + def partitions(self, state_name: str): + start = 0 + for batch in self.batches: + end = start + batch.request_count + if end > start: + payload = next((c.payload for c in batch.commits if c.state_name == state_name), None) + yield slice(start, end), payload + start = end diff --git a/lightllm/common/attention_state/manager.py b/lightllm/common/attention_state/manager.py index 6fb18c7ea1..835778a844 100644 --- a/lightllm/common/attention_state/manager.py +++ b/lightllm/common/attention_state/manager.py @@ -397,12 +397,19 @@ def commit_bounded_att_state( req_indices: torch.Tensor, accepted_steps: torch.Tensor, role: str, + *, + context=None, ) -> None: """传入 CUDA 请求索引和 int32 接受步数，包含基准输入步；无参与 State 时为空操作。""" assert req_indices.is_cuda and accepted_steps.is_cuda, "commit inputs must be CUDA tensors" assert role in MODEL_ROLES, f"unknown model role {role}" - for state in self.get_bounded_attention_states(role): - state.commit_bounded_att_state(req_indices, accepted_steps) + states = self.get_bounded_attention_states(role) + if context is not None: + context.validate(req_indices.numel(), {state.state_name for state in states}) + for state in states: + state.validate_bounded_att_commit(context) + for state in states: + state.commit_bounded_att_state(req_indices, accepted_steps, context=context) # ======================================================================== # 四、CPU/DP/PD 操作均由对应 Operator 提供。 diff --git a/lightllm/common/basemodel/attention/base_att.py b/lightllm/common/basemodel/attention/base_att.py index 28db5bf474..5e08a1d711 100644 --- a/lightllm/common/basemodel/attention/base_att.py +++ b/lightllm/common/basemodel/attention/base_att.py @@ -10,7 +10,7 @@ from lightllm.utils.envs_utils import get_env_start_args if TYPE_CHECKING: - from lightllm.common.linear_state import LinearStateCommit + from lightllm.common.attention_state.commit import StateCommit from lightllm.common.basemodel.basemodel import TpPartBaseModel from lightllm.common.basemodel.infer_struct import InferStateInfo @@ -162,8 +162,15 @@ def prefill_att( @dataclass class BaseDecodeAttState(BaseAttState): - def get_linear_state_commit(self) -> Optional["LinearStateCommit"]: - return None + def get_state_commits(self) -> tuple["StateCommit", ...]: + return () + + def clone_for_warmup(self, infer_state): + import copy + + state = copy.copy(self) + state._infer_state_weak_ref = weakref.ref(infer_state) + return state def copy_for_decode_cuda_graph(self, new_state: "BaseDecodeAttState"): """Copy tensor metadata into the buffers owned by the captured graph.""" diff --git a/lightllm/common/basemodel/attention/linear/gdn.py b/lightllm/common/basemodel/attention/linear/gdn.py index 32403a379f..2018f96031 100644 --- a/lightllm/common/basemodel/attention/linear/gdn.py +++ b/lightllm/common/basemodel/attention/linear/gdn.py @@ -1,29 +1,19 @@ from lightllm.common.attention_state.bounded_att import GDNState -from lightllm.common.linear_state import get_linear_state +from .state import LinearDecodeAttState, LinearPrefillAttState, get_resource import dataclasses import weakref from abc import ABC, abstractmethod -from typing import TYPE_CHECKING, Optional +from typing import TYPE_CHECKING import torch -from ..base_att import BaseAttBackend, BasePrefillAttState, BaseDecodeAttState, AttControl +from ..base_att import BaseAttBackend, AttControl from lightllm.utils.envs_utils import get_env_start_args from lightllm.common.basemodel.triton_kernel.linear_att.causal_conv1d import causal_conv1d_fn from lightllm.common.basemodel.triton_kernel.linear_att.fused_gdn_gating import fused_gdn_gating from lightllm.common.basemodel.triton_kernel.linear_att.gdn_decode_pack import conv_pack_gdn_decode_inputs -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.mtp_state_params import ( - build_dynamic_mtp_linear_att_state_params, -) -from lightllm.common.basemodel.triton_kernel.linear_att.fla.ops import ( - fused_recurrent_gated_delta_rule, -) if TYPE_CHECKING: - from lightllm.common.linear_state import LinearStateCommit, SSMRunConfig from lightllm.common.basemodel.basemodel import TpPartBaseModel from lightllm.common.basemodel.infer_struct import InferStateInfo from lightllm.models.qwen3next.infer_struct import Qwen3NextInferStateInfo @@ -112,18 +102,8 @@ def create_att_decode_state(self, infer_state: "InferStateInfo") -> "GDNLinearAt @dataclasses.dataclass -class GDNLinearAttPrefillAttState(BasePrefillAttState): - backend: GDNLinearAttBackend = None - b_conv_buffer_idx: torch.Tensor = None - b_ssm_buffer_idx: torch.Tensor = None - - def init_state(self): - # 每次 _prefill 都会在 runtime infer_state 上调用 init_state。 - # prefill cuda graph 回调必须走 new_infer_state.prefill_att_state1， - # 才能读到这里按当前 batch（含 token padding 后的 dummy request）更新的索引。 - linear_state_manager = get_linear_state(self.infer_state, GDNState) - self.b_conv_buffer_idx, self.b_ssm_buffer_idx = linear_state_manager.prepare_prefill(self.infer_state.b_req_idx) - return +class GDNLinearAttPrefillAttState(LinearPrefillAttState): + state_class = GDNState def prefill_att( self, @@ -142,14 +122,22 @@ def prefill_att( layer_num = linear_att_dict["layer_num"] backend: GDNLinearAttBackend = self.backend - conv_states, ssm_states = get_linear_state(self.infer_state, GDNState).get_layer_view(layer_num) + resource = get_resource(self, layer_num) + conv_states, ssm_states = resource.get_layer_view(layer_num) # 在开启了mtp的时候，conv 状态的最后一维可能存在冗余的部分，需要进行切片对齐。 # prefill 模式下，使用不到这几个维度，所以需要扣除掉， if backend.mtp_step > 0: conv_states = conv_states[:, :, : -backend.mtp_step] mixed_qkv, z, b, a = backend._split_qkvzba(mixed_qkvzba) core_attn_out = self._gdn_prefill_kernel( - mixed_qkv, conv_states, ssm_states, a, b, self.infer_state, layer_weight + mixed_qkv, + conv_states, + ssm_states, + a, + b, + self.infer_state, + layer_weight, + self.state_indices[resource.state_name], ) return core_attn_out, z @@ -162,17 +150,19 @@ def _gdn_prefill_kernel( b: torch.Tensor, infer_state: "Qwen3NextInferStateInfo", layer_weight: "Qwen3NextTransformerLayerWeight", + state_indices, ): g, beta = fused_gdn_gating(layer_weight.linear_A_log.weight, a, b, layer_weight.linear_dt_bias.weight) mixed_qkv = mixed_qkv.transpose(0, 1) + conv_indices, ssm_indices = state_indices backend: GDNLinearAttBackend = self.backend out_tensor = causal_conv1d_fn( mixed_qkv, layer_weight.linear_conv1d.mm_param.weight, bias=layer_weight.linear_conv1d.bias, query_start_loc=infer_state.b1_cu_q_seq_len, - cache_indices=self.b_conv_buffer_idx, + cache_indices=conv_indices, has_initial_state=infer_state.b_ready_cache_len > 0, conv_states=conv_states, activation=backend.activation, @@ -181,7 +171,7 @@ def _gdn_prefill_kernel( # Recurrent processing query, key, value = backend._rearrange_mixed_qkv(mixed_qkv) - initial_state = ssm_states[self.b_ssm_buffer_idx] + initial_state = ssm_states[ssm_indices] # g and beta have shape (total_tokens, num_heads), need to unsqueeze to get (1, total_tokens, num_heads) core_attn_out, last_recurrent_state = backend.prefill_kernel( q=query, @@ -197,68 +187,16 @@ def _gdn_prefill_kernel( # The chunk kernel accumulates the recurrent state in float32 even when # the state cache is configured as bfloat16. Advanced indexing does # not perform an implicit dtype conversion for index_put. - ssm_states[self.b_ssm_buffer_idx] = last_recurrent_state.to(ssm_states.dtype, copy=False) + ssm_states[ssm_indices] = last_recurrent_state.to(ssm_states.dtype, copy=False) return core_attn_out @dataclasses.dataclass -class GDNLinearAttDecodeAttState(BaseDecodeAttState): - run_config: Optional["SSMRunConfig"] = None - backend: GDNLinearAttBackend = None - 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 - - def get_linear_state_commit(self) -> Optional["LinearStateCommit"]: - return get_linear_state(self.infer_state, GDNState).create_commit(self.run_config) - - 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: - self._init_normal_decode_state() - elif self.backend.uses_dynamic_spec_verify_layout(): - self._init_dynamic_mtp_decode_state() - else: - self._init_fixed_mtp_decode_state(draft_step) - linear_state_manager = get_linear_state(self.infer_state, GDNState) - if draft_step == 0: - linear_state_manager.prepare_decode(self.b_conv_buffer_idx, single_token=True) - else: - ( - self.b_conv_buffer_idx, - self.b_ssm_buffer_idx, - self.b_num_accepted_tokens, - ) = linear_state_manager.prepare_mtp_decode(self.b_conv_buffer_idx, self.b1_mtp_cu_q_seq_len) - - def _init_normal_decode_state(self): - self.b_conv_buffer_idx = self.infer_state.b_req_idx - self.b_ssm_buffer_idx = self.infer_state.b_req_idx - - def _init_dynamic_mtp_decode_state(self): - (self.b1_mtp_cu_q_seq_len, self.b_conv_buffer_idx,) = build_dynamic_mtp_linear_att_state_params( - b_req_idx=self.infer_state.b_req_idx, - b_mtp_index=self.infer_state.b_mtp_index, - hold_req_id=self.infer_state.req_manager.HOLD_REQUEST_ID, - ) +class GDNLinearAttDecodeAttState(LinearDecodeAttState): + state_class = GDNState - def _init_fixed_mtp_decode_state(self, draft_step: int) -> None: - mtp_size = draft_step + 1 - batch_size = self.infer_state.batch_size - assert batch_size % mtp_size == 0, ( - "GDN fixed-layout decode requires batch_size to be divisible by draft_step + 1, " - f"got batch_size={batch_size}, draft_step={draft_step}." - ) - - att_batch_size = batch_size // mtp_size - self.b1_mtp_cu_q_seq_len = torch.arange( - 0, - batch_size + 1, - mtp_size, - dtype=torch.int32, - device=self.infer_state.b_req_idx.device, - ) - self.b_conv_buffer_idx = self.infer_state.b_req_idx.view(att_batch_size, mtp_size)[:, 0].contiguous() + def get_draft_step(self): + return self.backend.model.mtp_manager.get_decode_draft_step(self.backend.model.is_mtp_draft_model) def decode_att( self, @@ -278,14 +216,15 @@ def decode_att( backend: GDNLinearAttBackend = self.backend mixed_qkv, z, b, a = backend._split_qkvzba(mixed_qkvzba) - conv_states, ssm_states = get_linear_state(self.infer_state, GDNState).get_layer_view(layer_num) + forward = self.get_forward(layer_num) + conv_states, _ = forward.resource.get_layer_view(layer_num) draft_step = self.backend.model.mtp_manager.get_decode_draft_step(self.backend.model.is_mtp_draft_model) if draft_step > 0: core_attn_out = self._gdn_mtp_kernel( mixed_qkv, conv_states, - ssm_states, + forward, a, b, self.infer_state, @@ -297,7 +236,7 @@ def decode_att( mixed_qkv, z, conv_states, - ssm_states, + forward, a, b, self.infer_state, @@ -310,7 +249,7 @@ def _gdn_decode_kernel( mixed_qkv: torch.Tensor, z: torch.Tensor, conv_states: torch.Tensor, - ssm_states: torch.Tensor, + forward, a: torch.Tensor, b: torch.Tensor, infer_state: "Qwen3NextInferStateInfo", @@ -329,7 +268,7 @@ def _gdn_decode_kernel( conv_states, layer_weight.linear_conv1d.mm_param.weight, layer_weight.linear_conv1d.bias, - self.b_conv_buffer_idx, + forward.req_indices, backend.activation, backend.conv_kernel_dim, backend.tp_num_k_heads, @@ -337,36 +276,15 @@ def _gdn_decode_kernel( backend.tp_num_v_heads, backend.head_v_dim, ) - ssm_backend = get_linear_state(self.infer_state, GDNState).ssm_backend - if ssm_backend is not None: - linear_layer_idx = get_linear_state(self.infer_state, GDNState).get_storage_layer_index( - layer_weight.layer_num_ - ) - core_attn_out, self.run_config = ssm_backend.forward( - linear_layer_idx, - query, - key, - value, - a, - b, - layer_weight.linear_A_log.weight, - layer_weight.linear_dt_bias.weight, - self.b_conv_buffer_idx, - run_config=self.run_config, - ) - return core_attn_out, z - core_attn_out, _ = fused_recurrent_gated_delta_rule( - q=query, - k=key, - v=value, - initial_state=ssm_states, - inplace_final_state=True, - ssm_state_indices=self.b_ssm_buffer_idx, - use_qk_l2norm_in_kernel=True, - A_log=layer_weight.linear_A_log.weight, - dt_bias=layer_weight.linear_dt_bias.weight, - a_raw=a, - b_raw=b, + core_attn_out = forward.run_layer( + layer_weight.layer_num_, + query, + key, + value, + a, + b, + layer_weight.linear_A_log.weight, + layer_weight.linear_dt_bias.weight, ) return core_attn_out, z @@ -374,7 +292,7 @@ def _gdn_mtp_kernel( self, mixed_qkv: torch.Tensor, conv_states: torch.Tensor, - ssm_states: torch.Tensor, + forward, a: torch.Tensor, b: torch.Tensor, infer_state: "Qwen3NextInferStateInfo", @@ -386,7 +304,7 @@ def _gdn_mtp_kernel( backend: GDNLinearAttBackend = self.backend - cu_seqlens_q = self.b1_mtp_cu_q_seq_len + cu_seqlens_q = forward.cu_seqlens mixed_qkv = causal_conv1d_update_mtp( mixed_qkv, conv_states, @@ -394,47 +312,20 @@ def _gdn_mtp_kernel( mtp_step=backend.model.mtp_manager.get_decode_draft_step(backend.model.is_mtp_draft_model), bias=layer_weight.linear_conv1d.bias, activation=backend.activation, - conv_state_indices=self.b_conv_buffer_idx, - num_accepted_tokens=self.b_num_accepted_tokens, + conv_state_indices=forward.req_indices, + num_accepted_tokens=forward.accepted_token_counts, query_start_loc=cu_seqlens_q, ) query, key, value = backend._rearrange_mixed_qkv(mixed_qkv, decode=False) - ssm_backend = get_linear_state(self.infer_state, GDNState).ssm_backend - if ssm_backend is not None: - linear_layer_idx = get_linear_state(self.infer_state, GDNState).get_storage_layer_index( - layer_weight.layer_num_ - ) - core_attn_out, self.run_config = ssm_backend.forward( - linear_layer_idx, - query, - key, - value, - a, - b, - layer_weight.linear_A_log.weight, - layer_weight.linear_dt_bias.weight, - self.b_conv_buffer_idx, - cu_seqlens_q, - run_config=self.run_config, - ) - return core_attn_out - 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 - # .all() D2H sync stalled the GPU on the eager decode hot path; it is redundant here. - core_attn_out, _ = mtp_fused_recurrent_gated_delta_rule( - q=query, - k=key, - v=value, - initial_state=ssm_states, - cu_seqlens=cu_seqlens_q.to(torch.long), - ssm_state_indices=self.b_ssm_buffer_idx, - ssm_state_write_indices=self.b_ssm_buffer_idx, - num_accepted_tokens=self.b_num_accepted_tokens, - A_log=layer_weight.linear_A_log.weight, - dt_bias=layer_weight.linear_dt_bias.weight, - a_raw=a, - b_raw=b, + core_attn_out = forward.run_layer( + layer_weight.layer_num_, + query, + key, + value, + a, + b, + layer_weight.linear_A_log.weight, + layer_weight.linear_dt_bias.weight, ) return core_attn_out diff --git a/lightllm/common/basemodel/attention/linear/kda.py b/lightllm/common/basemodel/attention/linear/kda.py index 2fa1cbbba6..a985a6bc13 100644 --- a/lightllm/common/basemodel/attention/linear/kda.py +++ b/lightllm/common/basemodel/attention/linear/kda.py @@ -6,33 +6,26 @@ import dataclasses import weakref -from typing import TYPE_CHECKING, Optional +from typing import TYPE_CHECKING import torch from lightllm.common.attention_state.bounded_att import KDAState -from lightllm.common.linear_state import get_linear_state +from .state import LinearDecodeAttState, LinearPrefillAttState, get_resource from lightllm.common.basemodel.attention.base_att import ( AttControl, BaseAttBackend, - BaseDecodeAttState, - BasePrefillAttState, ) from lightllm.common.basemodel.triton_kernel.linear_att.causal_conv1d import ( causal_conv1d_fn, causal_conv1d_update, ) from lightllm.common.basemodel.triton_kernel.linear_att.fla.ops.kda import chunk_kda_with_fused_gate -from lightllm.common.basemodel.triton_kernel.linear_att.fla.ops.kda_decode import fused_recurrent_kda from lightllm.common.basemodel.triton_kernel.linear_att.causal_conv1d_mtp import ( causal_conv1d_update as causal_conv1d_update_mtp, ) -from lightllm.common.basemodel.triton_kernel.linear_att.mtp_state_params import ( - build_dynamic_mtp_linear_att_state_params, -) if TYPE_CHECKING: - from lightllm.common.linear_state import LinearStateCommit, SSMRunConfig from lightllm.common.basemodel.basemodel import TpPartBaseModel from lightllm.common.basemodel.infer_struct import InferStateInfo @@ -61,13 +54,8 @@ def split_qkv(self, mixed_qkv: torch.Tensor): @dataclasses.dataclass -class KDALinearAttPrefillAttState(BasePrefillAttState): - b_conv_buffer_idx: torch.Tensor = None - b_ssm_buffer_idx: torch.Tensor = None - - def init_state(self): - linear_state_manager = get_linear_state(self.infer_state, KDAState) - self.b_conv_buffer_idx, self.b_ssm_buffer_idx = linear_state_manager.prepare_prefill(self.infer_state.b_req_idx) +class KDALinearAttPrefillAttState(LinearPrefillAttState): + state_class = KDAState def prefill_att( self, @@ -86,12 +74,8 @@ def prefill_att( raw_beta = params["raw_beta"] backend: KDALinearAttBackend = self.backend - state = self.infer_state.att_state_manager.get_state( - model_role=self.infer_state.att_state_role, - state_type_filter="kda", - logical_layer_id=layer_num, - state_class=KDAState, - ) + state = get_resource(self, layer_num) + conv_indices, ssm_indices = self.state_indices[state.state_name] conv_states, ssm_states = state.get_layer_view(logical_layer_id=layer_num) conv_states = conv_states[..., : backend.conv_kernel_size - 1] mixed_qkv = causal_conv1d_fn( @@ -99,7 +83,7 @@ def prefill_att( layer_weight.get_merged_kda_conv_weight(), bias=None, query_start_loc=self.infer_state.b1_cu_q_seq_len, - cache_indices=self.b_conv_buffer_idx, + cache_indices=conv_indices, has_initial_state=self.infer_state.b_ready_cache_len > 0, conv_states=conv_states, activation="silu", @@ -112,7 +96,7 @@ def prefill_att( raw_gate = raw_gate.view(1, -1, backend.tp_hidden_size) raw_beta = raw_beta.view(1, -1, backend.tp_num_heads) - initial_state = ssm_states[self.b_ssm_buffer_idx].contiguous() + initial_state = ssm_states[ssm_indices].contiguous() output, final_state = chunk_kda_with_fused_gate( q=q, k=k, @@ -128,67 +112,13 @@ def prefill_att( safe_gate=True, lower_bound=backend.lower_bound, ) - ssm_states[self.b_ssm_buffer_idx] = final_state.to(ssm_states.dtype, copy=False) + ssm_states[ssm_indices] = final_state.to(ssm_states.dtype, copy=False) return output @dataclasses.dataclass -class KDALinearAttDecodeAttState(BaseDecodeAttState): - run_config: Optional["SSMRunConfig"] = None - 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 - - def get_linear_state_commit(self) -> Optional["LinearStateCommit"]: - return get_linear_state(self.infer_state, KDAState).create_commit(self.run_config) - - def init_state(self): - mtp_step = self.backend.mtp_step - if mtp_step == 0: - self._init_normal_decode_state() - elif self.backend.uses_dynamic_spec_verify_layout(): - self._init_dynamic_mtp_decode_state() - else: - self._init_fixed_mtp_decode_state(mtp_step) - linear_state_manager = get_linear_state(self.infer_state, KDAState) - if mtp_step == 0: - linear_state_manager.prepare_decode(self.b_conv_buffer_idx, single_token=True) - else: - ( - self.b_conv_buffer_idx, - self.b_ssm_buffer_idx, - self.b_num_accepted_tokens, - ) = linear_state_manager.prepare_mtp_decode(self.b_conv_buffer_idx, self.b1_mtp_cu_q_seq_len) - - def _init_normal_decode_state(self): - self.b_conv_buffer_idx = self.infer_state.b_req_idx - self.b_ssm_buffer_idx = self.infer_state.b_req_idx - - def _init_dynamic_mtp_decode_state(self): - (self.b1_mtp_cu_q_seq_len, self.b_conv_buffer_idx,) = build_dynamic_mtp_linear_att_state_params( - b_req_idx=self.infer_state.b_req_idx, - b_mtp_index=self.infer_state.b_mtp_index, - hold_req_id=self.infer_state.req_manager.HOLD_REQUEST_ID, - ) - - def _init_fixed_mtp_decode_state(self, mtp_step: int) -> None: - mtp_size = mtp_step + 1 - batch_size = self.infer_state.batch_size - assert batch_size % mtp_size == 0, ( - "KDA fixed-layout decode requires batch_size to be divisible by mtp_step + 1, " - f"got batch_size={batch_size}, mtp_step={mtp_step}." - ) - - att_batch_size = batch_size // mtp_size - self.b1_mtp_cu_q_seq_len = torch.arange( - 0, - batch_size + 1, - mtp_size, - dtype=torch.int32, - device=self.infer_state.b_req_idx.device, - ) - self.b_conv_buffer_idx = self.infer_state.b_req_idx.view(att_batch_size, mtp_size)[:, 0].contiguous() +class KDALinearAttDecodeAttState(LinearDecodeAttState): + state_class = KDAState def decode_att( self, @@ -207,21 +137,16 @@ def decode_att( raw_beta = params["raw_beta"] backend: KDALinearAttBackend = self.backend - state = self.infer_state.att_state_manager.get_state( - model_role=self.infer_state.att_state_role, - state_type_filter="kda", - logical_layer_id=layer_num, - state_class=KDAState, - ) - conv_states, ssm_states = state.get_layer_view(logical_layer_id=layer_num) - conv_kwargs = dict(bias=None, activation="silu", conv_state_indices=self.b_conv_buffer_idx) + forward = self.get_forward(layer_num) + conv_states, _ = forward.resource.get_layer_view(layer_num) + conv_kwargs = dict(bias=None, activation="silu", conv_state_indices=forward.req_indices) conv_update = causal_conv1d_update if backend.mtp_step > 0: conv_update = causal_conv1d_update_mtp conv_kwargs.update( mtp_step=backend.mtp_step, - num_accepted_tokens=self.b_num_accepted_tokens, - query_start_loc=self.b1_mtp_cu_q_seq_len, + num_accepted_tokens=forward.accepted_token_counts, + query_start_loc=forward.cu_seqlens, ) mixed_qkv = conv_update(mixed_qkv, conv_states, layer_weight.get_merged_kda_conv_weight(), **conv_kwargs) q, k, v = backend.split_qkv(mixed_qkv) @@ -231,36 +156,13 @@ def decode_att( v = v.view(*shape, backend.tp_num_heads, backend.head_dim) raw_gate = raw_gate.view(*shape, backend.tp_hidden_size) raw_beta = raw_beta.view(*shape, backend.tp_num_heads) - ssm_backend = get_linear_state(self.infer_state, KDAState).ssm_backend - if ssm_backend is not None: - linear_layer_idx = get_linear_state(self.infer_state, KDAState).get_storage_layer_index(layer_num) - output, self.run_config = ssm_backend.forward( - linear_layer_idx, - q, - k, - v, - raw_gate.view(-1, backend.tp_hidden_size), - raw_beta.view(-1, backend.tp_num_heads), - layer_weight.linear_A_log.weight, - layer_weight.linear_dt_bias.weight, - self.b_conv_buffer_idx, - self.b1_mtp_cu_q_seq_len, - run_config=self.run_config, - ) - return output - output, _ = fused_recurrent_kda( - q=q, - k=k, - v=v, - raw_gate=raw_gate, - raw_beta=raw_beta, - a_log=layer_weight.linear_A_log.weight, - gate_bias=layer_weight.linear_dt_bias.weight, - initial_state=ssm_states, - lower_bound=backend.lower_bound, - inplace_final_state=True, - ssm_state_indices=self.b_ssm_buffer_idx, - cu_seqlens=self.b1_mtp_cu_q_seq_len, - num_accepted_tokens=self.b_num_accepted_tokens, + return forward.run_layer( + layer_num, + q, + k, + v, + raw_gate.view(-1, backend.tp_hidden_size), + raw_beta.view(-1, backend.tp_num_heads), + layer_weight.linear_A_log.weight, + layer_weight.linear_dt_bias.weight, ) - return output diff --git a/lightllm/common/basemodel/attention/linear/state.py b/lightllm/common/basemodel/attention/linear/state.py new file mode 100644 index 0000000000..64d2c2243c --- /dev/null +++ b/lightllm/common/basemodel/attention/linear/state.py @@ -0,0 +1,89 @@ +import copy +from dataclasses import dataclass, field + +import torch + +from lightllm.common.basemodel.attention.base_att import BaseDecodeAttState, BasePrefillAttState +from lightllm.common.basemodel.triton_kernel.linear_att.mtp_state_params import ( + build_dynamic_mtp_linear_att_state_params, +) + + +def get_resources(att_state): + return tuple( + state + for state in att_state.infer_state.att_state_manager.get_bounded_attention_states( + att_state.infer_state.att_state_role + ) + if isinstance(state, att_state.state_class) + ) + + +def get_resource(att_state, logical_layer): + return att_state.infer_state.att_state_manager.get_state( + model_role=att_state.infer_state.att_state_role, + state_type_filter=att_state.state_class.algorithm, + logical_layer_id=logical_layer, + state_class=att_state.state_class, + ) + + +@dataclass +class LinearPrefillAttState(BasePrefillAttState): + state_indices: dict = field(default_factory=dict) + + def init_state(self): + + self.state_indices = { + state.state_name: state.prepare_prefill(self.infer_state.b_req_idx) for state in get_resources(self) + } + + +@dataclass +class LinearDecodeAttState(BaseDecodeAttState): + forwards: dict = field(default_factory=dict) + + def get_draft_step(self): + return self.backend.mtp_step + + def init_state(self): + draft_step = self.get_draft_step() + infer_state = self.infer_state + reqs, cu_seqlens = infer_state.b_req_idx, None + if draft_step: + if self.backend.uses_dynamic_spec_verify_layout(): + cu_seqlens, reqs = build_dynamic_mtp_linear_att_state_params( + b_req_idx=reqs, + b_mtp_index=infer_state.b_mtp_index, + hold_req_id=infer_state.req_manager.HOLD_REQUEST_ID, + ) + else: + width = draft_step + 1 + if infer_state.batch_size % width: + raise ValueError("fixed-layout decode requires batch_size divisible by draft_step + 1") + cu_seqlens = torch.arange( + 0, + infer_state.batch_size + 1, + width, + dtype=torch.int32, + device=reqs.device, + ) + reqs = reqs.view(-1, width)[:, 0].contiguous() + self.forwards = {state.state_name: state.prepare_decode(reqs, cu_seqlens) for state in get_resources(self)} + + def get_forward(self, logical_layer): + return self.forwards[get_resource(self, logical_layer).state_name] + + def get_state_commits(self): + return tuple(commit for forward in self.forwards.values() if (commit := forward.finish()) is not None) + + def clone_for_warmup(self, infer_state): + state = super().clone_for_warmup(infer_state) + state.forwards = {name: copy.copy(forward) for name, forward in self.forwards.items()} + return state + + def copy_for_decode_cuda_graph(self, new_state): + if self.forwards.keys() != new_state.forwards.keys(): + raise ValueError("captured linear attention resources changed") + for name, forward in self.forwards.items(): + forward.copy_inputs(new_state.forwards[name]) diff --git a/lightllm/common/basemodel/basemodel.py b/lightllm/common/basemodel/basemodel.py index af2ae932be..e8e112cc97 100755 --- a/lightllm/common/basemodel/basemodel.py +++ b/lightllm/common/basemodel/basemodel.py @@ -577,19 +577,18 @@ def _create_padded_prefill_model_input(self, model_input: ModelInput, new_handle return new_model_input def _create_model_output(self, post_output: PostLayerOutput, infer_state: InferStateInfo) -> ModelOutput: - linear_state_commit = None - for state in (infer_state.decode_att_state, infer_state.decode_att_state1): - if state is not None: - commit = state.get_linear_state_commit() - if commit is not None: - assert linear_state_commit is None, "multiple linear state commits in one forward" - linear_state_commit = commit + state_commits = tuple( + commit + for state in (infer_state.decode_att_state, infer_state.decode_att_state1) + if state is not None + for commit in state.get_state_commits() + ) output = ModelOutput( logits=post_output.logits.contiguous(), mtp_collector=infer_state.hidden_collector.finish_output(infer_state=infer_state), prompt_logics=infer_state.prompt_logics, logits_token_ids=post_output.logits_token_ids, - linear_state_commit=linear_state_commit, + state_commits=state_commits, ) return output diff --git a/lightllm/common/basemodel/batch_objs.py b/lightllm/common/basemodel/batch_objs.py index b2e2b54752..8e4b39688e 100644 --- a/lightllm/common/basemodel/batch_objs.py +++ b/lightllm/common/basemodel/batch_objs.py @@ -7,7 +7,7 @@ from lightllm.utils.tensor_utils import tensor_to_no_ref_tensor if TYPE_CHECKING: - from lightllm.common.linear_state import LinearStateCommit + from lightllm.common.attention_state.commit import StateCommit @dataclass @@ -206,7 +206,7 @@ class ModelOutput: logits_token_ids: Optional[torch.Tensor] = None # Detached state-commit metadata for this forward/graph; preserved when unpadding. - linear_state_commit: Optional["LinearStateCommit"] = None + state_commits: tuple["StateCommit", ...] = () def __post_init__(self) -> None: if self.mtp_collector is None: diff --git a/lightllm/common/basemodel/cuda_graph.py b/lightllm/common/basemodel/cuda_graph.py index 39de136d40..b85ce3b0f2 100644 --- a/lightllm/common/basemodel/cuda_graph.py +++ b/lightllm/common/basemodel/cuda_graph.py @@ -2,7 +2,6 @@ import copy import bisect import math -import weakref from typing import Optional import torch @@ -29,8 +28,7 @@ def _copy_warmup_state(infer_state: InferStateInfo) -> InferStateInfo: for name in ("decode_att_state", "decode_att_state1"): state = getattr(infer_state, name, None) if state is not None: - state = copy.copy(state) - state._infer_state_weak_ref = weakref.ref(warmup_state) + state = state.clone_for_warmup(warmup_state) setattr(warmup_state, name, state) return warmup_state diff --git a/lightllm/common/basemodel/triton_kernel/linear_att/compact_ssm.py b/lightllm/common/basemodel/triton_kernel/linear_att/compact_ssm.py index 7173a3eed6..c16a9d7983 100644 --- a/lightllm/common/basemodel/triton_kernel/linear_att/compact_ssm.py +++ b/lightllm/common/basemodel/triton_kernel/linear_att/compact_ssm.py @@ -5,14 +5,9 @@ accepted prefix with the same per-token rounding as the recurrent baseline. """ -import torch import triton import triton.language as tl -from lightllm.common.linear_state import SSMRunConfig - -from .ssm_autotune import resolve_run_config - @triton.jit def _compact( @@ -115,154 +110,3 @@ def _compact( state = state.to(State.dtype.element_ty).to(tl.float32) if COMMIT: tl.store(sp, state, (kk[:, None] < K) & (vv[None, :] < V)) - - -class CompactSSMBackend: - def __init__( - self, - state, - verify_width, - activation_dtype, - *, - num_key_heads=None, - kda=False, - gate_lower_bound=-5.0, - run_config=None, - ): - self.state = state - self.kda = kda - self.gate_lower_bound = gate_lower_bound - self.verify_width = verify_width - # Verify and accepted-state reconstruction must share the same layout. - self._config_is_fixed = run_config is not None - self.batch_configs = {} - self.run_config = dict(run_config or {"BV": 8, "num_warps": 1}) - assert self.run_config["BV"] in (8, 16, 32, 64, 128) - assert self.run_config["num_warps"] in (1, 2, 4, 8) - layers, slots, hv, k, v = state.shape - self.num_key_heads = hv if num_key_heads is None else num_key_heads - assert self.num_key_heads > 0 and hv % self.num_key_heads == 0 - self.hold_req_id = slots - 1 - shape = (layers, slots, hv, verify_width) - self.keys = torch.empty( - (layers, slots, self.num_key_heads, verify_width, 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 kda else shape, device=state.device, dtype=torch.float32) - self.betas = torch.empty(shape, device=state.device, dtype=torch.float32) - - def clear_history(self, req): - # Verify overwrites scratch before accepting updates; no history spans rounds. - pass - - def prepare_decode(self, reqs, cu_seqlens=None): - # Compact mode has no history spanning rounds. - return None - - def merge_accepted_updates(self, reqs): - # accept_updates already merged the accepted prefix into SSM state. - pass - - def get_accepted_state(self, req_idx): - return self.state[:, req_idx] - - def get_run_config(self, q, k, v, a, b, a_log, bias, cu_seqlens=None) -> SSMRunConfig: - config = resolve_run_config(self, "compact", q, k, v, a, b, a_log, bias, cu_seqlens) - return SSMRunConfig(**config) - - def forward(self, linear_layer_idx, q, k, v, a, b, a_log, bias, reqs, cu_seqlens=None, run_config=None): - assert cu_seqlens is not None, "compact replay is only used for speculative verify" - config = run_config or self.get_run_config(q, k, v, a, b, a_log, bias, cu_seqlens) - if not isinstance(config, SSMRunConfig): - config = SSMRunConfig(**config) - _, slots, hv, kd, vd = self.state.shape - assert self.num_key_heads in (q.shape[-2], hv) - out = torch.empty_like(v) - # Match native reduction layouts: changing BV can alter BF16 rounding. - bv = config.BV - _compact[(triton.cdiv(vd, bv), reqs.numel(), hv)]( - q, - k, - v, - a, - b, - a_log, - bias, - self.state[linear_layer_idx], - self.keys[linear_layer_idx], - self.values[linear_layer_idx], - self.decays[linear_layer_idx], - self.betas[linear_layer_idx], - 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, - self.num_key_heads, - kd, - vd, - self.verify_width, - self.hold_req_id, - triton.next_power_of_2(kd), - bv, - False, - self.kda, - self.gate_lower_bound, - num_warps=config.num_warps, - num_stages=config.num_stages, - ) - return out, config - - def accept_updates(self, reqs, accepted_state_indices, run_config=None): - """Merge this round's prefix; accepted_state_indices contains per-request last accepted indexes.""" - assert run_config is not None or not self.batch_configs, "compact acceptance requires the verify configuration" - layers, slots, hv, kd, vd = self.state.shape - config = run_config or self.run_config - if not isinstance(config, SSMRunConfig): - config = SSMRunConfig(**config) - bv = config.BV - _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_state_indices, - None, - 0, - 0, - 0, - 0, - 0, - slots, - hv, - hv, - self.num_key_heads, - kd, - vd, - self.verify_width, - self.hold_req_id, - triton.next_power_of_2(kd), - bv, - True, - self.kda, - self.gate_lower_bound, - num_warps=config.num_warps, - num_stages=config.num_stages, - ) diff --git a/lightllm/common/basemodel/triton_kernel/linear_att/replayssm.py b/lightllm/common/basemodel/triton_kernel/linear_att/replayssm.py index fbedde9202..c50e0c680f 100644 --- a/lightllm/common/basemodel/triton_kernel/linear_att/replayssm.py +++ b/lightllm/common/basemodel/triton_kernel/linear_att/replayssm.py @@ -9,14 +9,9 @@ Algorithm reference: https://dao-lab.ai/blog/2026/replayssm/ """ -import torch import triton import triton.language as tl -from lightllm.common.linear_state import SSMRunConfig - -from .ssm_autotune import resolve_run_config - @triton.jit def _accept_updates( @@ -585,263 +580,3 @@ def _materialize( out_slot = linear_layer_idx * HV + head if SNAPSHOT else slot out = Output + out_slot * K * V + kk[:, None] * V + vv[None, :] tl.store(out, state, (kk[:, None] < K) & (vv[None, :] < V)) - - -class ReplaySSMBackend: - """Request-owned scratch; canonical CPU/PD checkpoints never include it.""" - - def __init__( - self, - state, - capacity, - verify_width, - activation_dtype=torch.bfloat16, - *, - num_key_heads=None, - kda=False, - gate_lower_bound=-5.0, - compute_mode="inline", - run_config=None, - fold_method="raw", - ): - assert state.dtype in (torch.float32, torch.bfloat16), "ReplaySSM requires FP32 or BF16 SSM state" - assert verify_width > 0 and capacity >= max(4, verify_width) and capacity & (capacity - 1) == 0 - assert compute_mode in ("inline", "precompute", "parallel") - assert fold_method in ("raw", "derived") - assert fold_method == "raw" or compute_mode == "parallel" - if compute_mode == "parallel": - assert not kda and verify_width > 1, "Parallel ReplaySSM requires GDN MTP" - assert state.shape[-2:] == (128, 128), "Parallel ReplaySSM requires 128x128 heads" - assert activation_dtype == torch.bfloat16, "Parallel ReplaySSM requires BF16 activations" - assert state.dtype == torch.bfloat16, "Parallel ReplaySSM currently supports BF16 state" - self.fold_method = fold_method - self.state = state - self.kda = kda - self.gate_lower_bound = gate_lower_bound - self.fold_programs = ( - 8 * torch.cuda.get_device_properties(state.device).multi_processor_count if state.is_cuda else 0 - ) - self.capacity = capacity - self.verify_width = verify_width - self.compute_mode = compute_mode - self.run_config = run_config - self._config_is_fixed = run_config is not None - self.batch_configs = {} - layers, slots, hv, k, v = state.shape - # Omission preserves the redundant layout for existing callers and A/B. - self.num_key_heads = hv if num_key_heads is None else num_key_heads - assert self.num_key_heads > 0 and hv % self.num_key_heads == 0 - self.hold_req_id = slots - 1 - self.cursors = torch.zeros(slots, dtype=torch.int32, device=state.device) - self.forward_cursors = torch.zeros_like(self.cursors) - # Alternate halves on fold so CTAs cannot overwrite history another - # V tile is still reading. Cursor packs phase (2*L) and count (0..L). - # Derived corrections can exceed the activation/state range even when - # the recurrence stays finite. Only BF16 has enough exponent range for - # the short-history narrowing optimization; FP16 must retain FP32. - history_dtype = ( - activation_dtype - if activation_dtype == torch.bfloat16 - and verify_width > 1 - and compute_mode == "precompute" - and capacity <= 8 - else torch.float32 - ) - key_shape = (layers, slots, self.num_key_heads, 2 * capacity, k) - self.keys = torch.empty(key_shape, dtype=history_dtype, device=state.device) - self.deltas = torch.empty((layers, slots, hv, 2 * capacity, v), dtype=history_dtype, device=state.device) - self.gates = torch.empty( - (layers, slots, hv, 2 * capacity, k) if kda else (layers, slots, hv, 2 * capacity), - dtype=torch.float32, - device=state.device, - ) - self.raw_keys = torch.empty(key_shape, 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 clear_history(self, req): - """Invalidate history for a reused/restored request; leave its SSM state intact.""" - if type(req) is int: - # Scalar assignment copies a CPU tensor into the CUDA slot. - self.cursors[req].zero_() - else: - self.cursors[req] = 0 - - def prepare_decode(self, reqs, cu_seqlens=None): - """Snapshot history for all layers and prepare the base for accepted updates.""" - positions = torch.empty_like(reqs) - layers, _, hv, kd, vd = self.state.shape - active = torch.empty((reqs.numel() + 1,), dtype=torch.int32, device=reqs.device) - _prepare_fold[(1,)]( - reqs, - self.cursors, - positions, - self.forward_cursors, - cu_seqlens, - active, - reqs.numel(), - self.hold_req_id, - self.capacity, - cu_seqlens is not None, - triton.next_power_of_2(reqs.numel()), - self.verify_width, - self.compute_mode == "parallel", - num_warps=4, - ) - programs = min(self.fold_programs, triton.cdiv(vd, 32) * reqs.numel() * layers * hv) - _fold_history[(programs,)]( - self.state, - self.raw_keys, - self.raw_values, - self.gates, - self.betas, - reqs, - positions, - active, - reqs.numel(), - layers, - hv, - self.num_key_heads, - kd, - vd, - self.capacity, - triton.next_power_of_2(kd), - 32, - self.state.stride(0), - self.raw_keys.stride(0), - self.raw_values.stride(0), - self.gates.stride(0), - self.betas.stride(0), - self.kda, - self.keys, - self.deltas, - self.fold_method == "derived", - BL=max(16, self.capacity), - num_warps=2, - num_stages=1, - ) - return positions - - def accept_updates(self, reqs, accepted_state_indices=None): - """Accept this round's prefix; accepted_state_indices contains per-request last accepted indexes.""" - _accept_updates[(triton.cdiv(reqs.numel(), 256),)]( - self.cursors, - reqs, - accepted_state_indices, - reqs.numel(), - self.hold_req_id, - self.verify_width, - 256, - ) - - def merge_accepted_updates(self, reqs): - """Fold accepted history into SSM state and clear its history positions.""" - if reqs.numel() == 0: - return - self._materialize_accepted_state(reqs, self.state, snapshot=False) - self.cursors.index_fill_(0, reqs.long(), 0) - - def _materialize_accepted_state(self, reqs, output, snapshot): - layers, slots, hv, k, v = self.state.shape - config = self.run_config or {"BV": 32, "num_warps": 4} - bv = config["BV"] - _materialize[(triton.cdiv(v, bv), reqs.numel(), layers * hv)]( - self.state, - output, - self.raw_keys, - self.raw_values, - self.gates, - self.betas, - self.cursors, - reqs, - slots, - hv, - self.num_key_heads, - k, - v, - self.capacity, - self.hold_req_id, - triton.next_power_of_2(k), - bv, - snapshot, - self.kda, - num_warps=config["num_warps"], - num_stages=config.get("num_stages", 3), - ) - - def get_accepted_state(self, req_idx): - """Return a canonical checkpoint without changing the active state or history.""" - layers, _, hv, k, v = self.state.shape - output = torch.empty((layers, hv, k, v), dtype=self.state.dtype, device=self.state.device) - reqs = torch.tensor([req_idx], dtype=torch.int32, device=self.state.device) - self._materialize_accepted_state(reqs, output, snapshot=True) - return output - - def get_run_config(self, q, k, v, a, b, a_log, bias, cu_seqlens=None) -> SSMRunConfig: - if self.compute_mode == "parallel": - config = self.run_config or {"BV": 128, "num_warps": 2, "num_stages": 1} - else: - config = resolve_run_config(self, "replay", q, k, v, a, b, a_log, bias, cu_seqlens) - return SSMRunConfig(**config) - - def forward(self, linear_layer_idx, q, k, v, a, b, a_log, bias, reqs, cu_seqlens=None, run_config=None): - config = run_config or self.get_run_config(q, k, v, a, b, a_log, bias, cu_seqlens) - if not isinstance(config, SSMRunConfig): - config = SSMRunConfig(**config) - hv, kd, vd = self.state.shape[-3:] - assert self.num_key_heads in (q.shape[-2], hv) - axis = 1 if cu_seqlens is not None else 0 - out = torch.empty_like(v) - bv = config.BV - precompute_projections = self.verify_width > 1 and bv >= 16 and self.compute_mode in ("precompute", "parallel") - if self.compute_mode == "parallel": - assert precompute_projections, "Parallel ReplaySSM requires BV >= 16" - # Resolve tile dimensions in Python: bare max() in Triton JIT code - # lowers to a tensor operation, which cannot define static shapes. - bw = triton.next_power_of_2(self.verify_width) - _replay[(triton.cdiv(vd, bv), reqs.numel(), hv)]( - q, - k, - v, - a, - b, - a_log, - bias, - self.state[linear_layer_idx], - self.keys[linear_layer_idx], - self.deltas[linear_layer_idx], - self.gates[linear_layer_idx], - self.raw_keys[linear_layer_idx], - self.raw_values[linear_layer_idx], - self.betas[linear_layer_idx], - reqs, - self.forward_cursors, - cu_seqlens, - out, - q.stride(axis), - k.stride(axis), - v.stride(axis), - a.stride(0), - b.stride(0), - q.shape[-2], - hv, - self.num_key_heads, - kd, - vd, - self.capacity, - self.hold_req_id, - triton.next_power_of_2(kd), - bv, - cu_seqlens is not None, - precompute_projections, - bw, - reqs.numel() > 12 and self.compute_mode != "parallel", - self.kda, - self.gate_lower_bound, - self.compute_mode == "parallel", - BQK=max(16, 2 * bw), - BR=triton.next_power_of_2(max(16, self.capacity + bw)), - num_warps=config.num_warps, - num_stages=config.num_stages, - ) - return out, config diff --git a/lightllm/common/linear_state.py b/lightllm/common/linear_state.py deleted file mode 100644 index d39fb40813..0000000000 --- a/lightllm/common/linear_state.py +++ /dev/null @@ -1,246 +0,0 @@ -"""Worker-local recurrent state shared by attention, checkpoints, and PD I/O. - -The runner orders operations on its current CUDA stream. This component does -not allocate request IDs, choose prefix boundaries, or synchronize streams. -Attention backends retain responsibility for recurrent computation. -""" - -from dataclasses import dataclass -from typing import TYPE_CHECKING, NamedTuple, Optional, Sequence, Tuple, Union - -import torch - - -if TYPE_CHECKING: - from lightllm.common.basemodel.triton_kernel.linear_att.replayssm import ReplaySSMBackend - from lightllm.common.basemodel.triton_kernel.linear_att.compact_ssm import CompactSSMBackend - - -@dataclass(frozen=True, slots=True) -class SSMRunConfig: - BV: int - num_warps: int - num_stages: int = 3 - - def __post_init__(self): - for value in (self.BV, self.num_warps, self.num_stages): - if type(value) is not int: - raise TypeError("SSM kernel configuration must be an integer") - - -@dataclass(frozen=True, slots=True) -class LinearStateCommit: - _run_config: SSMRunConfig - - -class LinearStateTensors(NamedTuple): - conv: torch.Tensor - ssm: torch.Tensor - - -class LinearStateManager: - """Shared Conv/SSM lifecycle for GDN and KDA attention resources.""" - - def __init__( - self, - config, - request_capacity, - mtp_step=0, - device="cuda", - *, - ssm_state_mode="native", - ssm_replay_config=None, - kda=False, - gate_lower_bound=-5.0, - ): - self.config = config - self.request_capacity = request_capacity - self.mtp_step = mtp_step - self.ssm_slots_per_req = mtp_step + 1 if ssm_state_mode == "native" else 1 - self.mtp_state_index = torch.zeros(request_capacity, dtype=torch.int32, device=device) if mtp_step else None - self._logical_to_storage_layer = {layer: index for index, layer in enumerate(config.logical_layer_ids)} - layers = len(config.logical_layer_ids) - self.conv_state = torch.zeros( - (layers, request_capacity, config.conv_channels, config.conv_width + mtp_step), - dtype=config.conv_dtype, - device=device, - ) - self.ssm_state = torch.zeros( - (layers, request_capacity * self.ssm_slots_per_req, config.num_v_heads, config.key_dim, config.value_dim), - dtype=config.ssm_dtype, - device=device, - ) - ssm_replay_config = ssm_replay_config or {} - self.ssm_backend: Optional[Union["ReplaySSMBackend", "CompactSSMBackend"]] = None - if ssm_state_mode == "replay": - from lightllm.common.basemodel.triton_kernel.linear_att import replayssm - - self.ssm_backend = replayssm.ReplaySSMBackend( - self.ssm_state, - ssm_replay_config.get("cache_len", 16), - self.mtp_step + 1, - config.conv_dtype, - num_key_heads=config.num_k_heads, - kda=kda, - gate_lower_bound=gate_lower_bound, - compute_mode=ssm_replay_config.get("compute_mode", "inline"), - fold_method=ssm_replay_config.get("fold_method", "raw"), - ) - elif ssm_state_mode == "compact": - from lightllm.common.basemodel.triton_kernel.linear_att import compact_ssm - - self.ssm_backend = compact_ssm.CompactSSMBackend( - self.ssm_state, - self.mtp_step + 1, - config.conv_dtype, - num_key_heads=config.num_k_heads, - kda=kda, - gate_lower_bound=gate_lower_bound, - ) - - def _reset_tracking(self, req_idx): - if self.ssm_backend is not None: - self.ssm_backend.clear_history(req_idx) - if self.mtp_state_index is not None: - self.mtp_state_index[req_idx].zero_() - - def reset(self, req_idx): - """Initialize a newly allocated/reused request, including speculative slots.""" - self._reset_tracking(req_idx) - self.conv_state[:, req_idx].zero_() - start = req_idx * self.ssm_slots_per_req - self.ssm_state[:, start : start + self.ssm_slots_per_req].zero_() - - def prepare_prefill(self, reqs): - """Merge accepted history and return convolution and SSM buffer indexes.""" - if self.ssm_backend is not None: - self.ssm_backend.merge_accepted_updates(reqs) - return reqs, reqs * self.ssm_slots_per_req - - def prepare_decode(self, reqs, cu_seqlens=None, *, single_token=False): - if self.ssm_backend is not None: - self.ssm_backend.prepare_decode(reqs, cu_seqlens) - if single_token: - # Preserve the forward cursor snapshot before accepting the - # guaranteed token, including when forward is graph-replayed. - self.ssm_backend.accept_updates(reqs) - - def prepare_mtp_decode(self, reqs, cu_seqlens): - """Return Conv/SSM indexes and prior accepted token counts; prepare replay history.""" - accepted_token_counts = self.mtp_state_index[reqs] + 1 - accepted_token_counts = torch.where(reqs == self.mtp_state_index.numel() - 1, 1, accepted_token_counts) - if self.ssm_backend is not None: - ssm_indexes = reqs - else: - offsets = torch.arange(self.ssm_slots_per_req, device=reqs.device, dtype=reqs.dtype) - ssm_indexes = reqs[:, None] * self.ssm_slots_per_req + offsets[None, :] - self.prepare_decode(reqs, cu_seqlens) - return reqs, ssm_indexes, accepted_token_counts - - def create_commit(self, run_config: Optional[SSMRunConfig]) -> Optional[LinearStateCommit]: - from lightllm.common.basemodel.triton_kernel.linear_att.compact_ssm import CompactSSMBackend - - if run_config is not None and isinstance(self.ssm_backend, CompactSSMBackend): - return LinearStateCommit(run_config) - return None - - def commit_accepted_state( - self, - b_req_mtp_start_loc, - b_req_idx, - b_mtp_index, - accepted_index, - verify_width, - linear_state_batches: Optional[Sequence[Tuple[int, Optional[LinearStateCommit]]]] = None, - ): - """Commit accepted prefixes without requiring full-state materialization. - - ``linear_state_batches`` partitions logical requests in the same order as the MTP - start locations. Each commit belongs to the forward that wrote that - partition's scratch; padding rows do not participate in acceptance. - """ - from lightllm.common.basemodel.triton_kernel.mtp_utils import linear_att_mtp_state_index_update - from lightllm.common.basemodel.triton_kernel.linear_att import replayssm - - replay_cursors = None - if isinstance(self.ssm_backend, replayssm.ReplaySSMBackend) and self.ssm_backend.verify_width > 1: - replay_cursors = self.ssm_backend.cursors - linear_att_mtp_state_index_update( - req_to_mtp_state_index=self.mtp_state_index, - b_req_mtp_start_loc=b_req_mtp_start_loc, - b_req_idx=b_req_idx, - b_mtp_index=b_mtp_index, - accepted_index=accepted_index, - verify_width=verify_width, - replay_cursors=replay_cursors, - ) - if self.ssm_backend is None or replay_cursors is not None: - return - reqs = b_req_idx[b_req_mtp_start_loc.long()] - if linear_state_batches is None or isinstance(self.ssm_backend, replayssm.ReplaySSMBackend): - self.ssm_backend.accept_updates(reqs, self.mtp_state_index) - return - assert sum(size for size, _ in linear_state_batches) == reqs.numel() - start = 0 - for size, commit in linear_state_batches: - if size: - self.ssm_backend.accept_updates( - reqs[start : start + size], - self.mtp_state_index, - run_config=commit._run_config if commit is not None else None, - ) - start += size - - def get_accepted_state(self, req_idx): - """Return the current accepted prefix without mutating active state. - - The caller associates this prefix with its cache/transport token boundary. - These tensors may alias the active pool: enqueue consuming copies on the - current stream before advancing/releasing the request. Cross-stream or - cross-process callers must use their existing completion events before - recycling sources or destinations. No CPU transfer or sync is implicit. - """ - width = self.config.conv_width - conv = self.conv_state[:, req_idx, ..., :width] - if self.mtp_state_index is not None: - offsets = torch.arange(width, device=conv.device) + self.mtp_state_index[req_idx] - conv = self.conv_state[:, req_idx].index_select(-1, offsets) - if self.ssm_backend is not None: - ssm = self.ssm_backend.get_accepted_state(req_idx) - elif self.ssm_slots_per_req > 1: - index = req_idx * self.ssm_slots_per_req + self.mtp_state_index[req_idx] - ssm = self.ssm_state.index_select(1, index.long().view(1))[:, 0] - else: - ssm = self.ssm_state[:, req_idx] - return LinearStateTensors(conv, ssm) - - def prepare_restore(self, req_idx): - """Invalidate old history and return canonical destinations for restore. - - The caller must fill both tensors before the next forward. This lets PD - transport unpack TP shards directly into the pool without a staging copy. - The same stream/completion rules as ``get_accepted_state`` apply. - """ - self._reset_tracking(req_idx) - width = self.config.conv_width - return LinearStateTensors( - self.conv_state[:, req_idx, ..., :width], - self.ssm_state[:, req_idx * self.ssm_slots_per_req], - ) - - def restore(self, req_idx, conv, ssm): - target = self.prepare_restore(req_idx) - target.conv.copy_(conv, non_blocking=True) - target.ssm.copy_(ssm, non_blocking=True) - - def get_storage_layer_index(self, logical_layer_id): - return self._logical_to_storage_layer[logical_layer_id] - - -def get_linear_state(infer_state, state_class): - (state,) = ( - state - for state in infer_state.att_state_manager.get_bounded_attention_states(infer_state.att_state_role) - if isinstance(state, state_class) - ) - return state diff --git a/lightllm/common/ssm_types.py b/lightllm/common/ssm_types.py new file mode 100644 index 0000000000..636f7ceffe --- /dev/null +++ b/lightllm/common/ssm_types.py @@ -0,0 +1,13 @@ +from dataclasses import dataclass + + +@dataclass(frozen=True, slots=True) +class SSMRunConfig: + BV: int + num_warps: int + num_stages: int = 3 + + def __post_init__(self): + for value in (self.BV, self.num_warps, self.num_stages): + if type(value) is not int: + raise TypeError("SSM kernel configuration must be an integer") diff --git a/lightllm/server/router/model_infer/mode_backend/chunked_prefill/impl.py b/lightllm/server/router/model_infer/mode_backend/chunked_prefill/impl.py index 64174a4fef..b18d0a6186 100644 --- a/lightllm/server/router/model_infer/mode_backend/chunked_prefill/impl.py +++ b/lightllm/server/router/model_infer/mode_backend/chunked_prefill/impl.py @@ -1,3 +1,4 @@ +from lightllm.common.attention_state.commit import StateCommitBatch import time from typing import List @@ -298,7 +299,7 @@ def decode_mtp( b_req_idx=model_input.b_req_idx, b_req_mtp_start_loc=b_req_mtp_start_loc, b_mtp_index=model_input.b_mtp_index, - linear_state_batches=((req_num, model_output.linear_state_commit),), + state_commit_batches=(StateCommitBatch(req_num, model_output.state_commits),), ) accepted_index_cpu = g_pin_mem_manager.async_copy_from_gpu_tensor( key="accepted_index", diff --git a/lightllm/server/router/model_infer/mode_backend/dp_backend/impl.py b/lightllm/server/router/model_infer/mode_backend/dp_backend/impl.py index f7276e6f0e..00b298d463 100644 --- a/lightllm/server/router/model_infer/mode_backend/dp_backend/impl.py +++ b/lightllm/server/router/model_infer/mode_backend/dp_backend/impl.py @@ -1,3 +1,4 @@ +from lightllm.common.attention_state.commit import StateCommitBatch import time from typing import List, Tuple @@ -533,7 +534,7 @@ def decode_mtp(self, event_pack: OverlapEventPack, decode_reqs: List[InferReq]): b_req_idx=model_input.b_req_idx, b_req_mtp_start_loc=b_req_mtp_start_loc, b_mtp_index=model_input.b_mtp_index, - linear_state_batches=((req_num, model_output.linear_state_commit),), + state_commit_batches=(StateCommitBatch(req_num, model_output.state_commits),), ) accepted_index_cpu = g_pin_mem_manager.async_copy_from_gpu_tensor( key="accepted_index", @@ -808,9 +809,9 @@ def decode_overlap_mtp(self, event_pack: OverlapEventPack, decode_reqs: List[Inf b_req_idx=b_req_idx, b_req_mtp_start_loc=b_req_mtp_start_loc, b_mtp_index=b_mtp_index, - linear_state_batches=( - (real_request_num0, model_output0.linear_state_commit), - (real_request_num1, model_output1.linear_state_commit), + state_commit_batches=( + StateCommitBatch(real_request_num0, model_output0.state_commits), + StateCommitBatch(real_request_num1, model_output1.state_commits), ), ) mtp_accept_len0 = mtp_accept_len[:real_request_num0] diff --git a/lightllm/server/router/model_infer/mtp_speculative/utils.py b/lightllm/server/router/model_infer/mtp_speculative/utils.py index 98fe1ea809..3338f342f7 100644 --- a/lightllm/server/router/model_infer/mtp_speculative/utils.py +++ b/lightllm/server/router/model_infer/mtp_speculative/utils.py @@ -5,7 +5,7 @@ import torch -from lightllm.common.linear_state import LinearStateManager +from lightllm.common.attention_state.commit import MtpCommitContext from lightllm.common.basemodel.triton_kernel.mtp_utils import ( mtp_scatter_next_token_ids, @@ -24,7 +24,7 @@ def verify_mtp_tokens( b_req_idx: torch.Tensor, b_req_mtp_start_loc: torch.Tensor, b_mtp_index: torch.Tensor, - linear_state_batches=None, + state_commit_batches=(), ) -> Tuple[torch.Tensor, torch.Tensor]: """Verify target tokens and update recurrent MTP state when required.""" @@ -34,18 +34,20 @@ def verify_mtp_tokens( new_next_token_ids=next_token_ids, b_req_idx=b_req_idx, ) - for state in backend.model.att_state_manager.get_bounded_attention_states("main"): - if isinstance(state, LinearStateManager): - state.commit_accepted_state( - b_req_mtp_start_loc=b_req_mtp_start_loc, - b_req_idx=b_req_idx, - b_mtp_index=b_mtp_index, - accepted_index=accepted_index, - verify_width=backend.max_draft_step + 1, - linear_state_batches=linear_state_batches, - ) - else: - state.commit_bounded_att_state(b_req_idx[b_req_mtp_start_loc], accept_lengths) + context = MtpCommitContext( + b_req_mtp_start_loc=b_req_mtp_start_loc, + b_req_idx=b_req_idx, + b_mtp_index=b_mtp_index, + accepted_index=accepted_index, + verify_width=backend.max_draft_step + 1, + batches=state_commit_batches, + ) + backend.model.att_state_manager.commit_bounded_att_state( + b_req_idx[b_req_mtp_start_loc.long()], + accept_lengths, + role="main", + context=context, + ) return accept_lengths, accepted_index diff --git a/test/benchmark/kernels/benchmark_kda_ssm.py b/test/benchmark/kernels/benchmark_kda_ssm.py index 8a8c4cfe58..2ce999257e 100644 --- a/test/benchmark/kernels/benchmark_kda_ssm.py +++ b/test/benchmark/kernels/benchmark_kda_ssm.py @@ -14,8 +14,8 @@ import triton from lightllm.common.basemodel.triton_kernel.linear_att.fla.ops.kda_decode import fused_recurrent_kda -from lightllm.common.basemodel.triton_kernel.linear_att.replayssm import ReplaySSMBackend -from lightllm.common.basemodel.triton_kernel.linear_att.compact_ssm import CompactSSMBackend +from lightllm.common.attention_state.bounded_att.ssm.replay import ReplaySSMState +from lightllm.common.attention_state.bounded_att.ssm.compact import CompactSSMState from lightllm.common.basemodel.triton_kernel.linear_att.ssm_autotune import get_configs @@ -37,9 +37,9 @@ def workload(args, batch, width, dtype, mode, acceptance): state = torch.randn(layers, slots, heads, dim, dim, dtype=dtype, device="cuda") * 0.01 ssm_backend = None if mode == "compact": - ssm_backend = CompactSSMBackend(state, width, q.dtype, kda=True) + ssm_backend = CompactSSMState(state, width, q.dtype, kda=True) elif mode != "native": - ssm_backend = ReplaySSMBackend(state, 8, width, q.dtype, kda=True, compute_mode=mode) + ssm_backend = ReplaySSMState(state, 8, width, q.dtype, kda=True, compute_mode=mode) if ssm_backend is not None: ssm_backend._config_is_fixed = True indexes = reqs[:, None] * width + torch.arange(width, device="cuda", dtype=torch.int32) @@ -76,7 +76,7 @@ def cycle(): ssm_backend.forward(layer, q, k, v, gate, beta, alog, bias, reqs, cu) if ssm_backend is not None: ssm_backend.accept_updates(reqs, accepted) - if isinstance(ssm_backend, ReplaySSMBackend): + if isinstance(ssm_backend, ReplaySSMState): ssm_backend._materialize_accepted_state(reqs[:1], snapshot, snapshot=True) ssm_backend.merge_accepted_updates(reqs) else: diff --git a/unit_tests/common/att_resources/test_hybrid_layer_mapping.py b/unit_tests/common/att_resources/test_hybrid_layer_mapping.py index f93a87c1e9..4580453297 100644 --- a/unit_tests/common/att_resources/test_hybrid_layer_mapping.py +++ b/unit_tests/common/att_resources/test_hybrid_layer_mapping.py @@ -79,9 +79,9 @@ def test_hybrid_layer_mapping_and_mtp_states( device="meta", ) assert bounded_state.ssm_slots_per_req == (3 if mode == "native" else 1) - assert (bounded_state.ssm_backend is None) == (mode == "native") + assert bounded_state._ssm.state is bounded_state.ssm_state if mode == "replay": - assert bounded_state.ssm_backend.capacity == 8 + assert bounded_state._ssm.capacity == 8 manager = AttentionStateManager(request_capacity=1) manager.register_state(model_role="main", state=full_state) manager.register_state(model_role="main", state=bounded_state) diff --git a/unit_tests/common/att_resources/test_mtp_state_commit.py b/unit_tests/common/att_resources/test_mtp_state_commit.py index c8fe301217..5d4d9b8c23 100644 --- a/unit_tests/common/att_resources/test_mtp_state_commit.py +++ b/unit_tests/common/att_resources/test_mtp_state_commit.py @@ -110,8 +110,9 @@ def test_verify_commits_main_states_for_fixed_and_dynamic_widths(widths: tuple[i assert manager.draft0_gdn_conv_fp32_ssm_fp32_state.mtp_state_index.tolist() == [0, 0, 0] -def test_verify_routes_linear_commit_context_and_other_bounded_states(monkeypatch): +def test_verify_routes_one_generic_commit_context(monkeypatch): from lightllm.server.router.model_infer.mtp_speculative import utils + from lightllm.common.attention_state.commit import StateCommitBatch reqs = torch.tensor([2, 2, 0, 0], dtype=torch.int32) starts = torch.tensor([0, 2], dtype=torch.int32) @@ -120,10 +121,7 @@ def test_verify_routes_linear_commit_context_and_other_bounded_states(monkeypatc mask = torch.tensor([1, 1, 1, 0], dtype=torch.int32) monkeypatch.setattr(utils, "mtp_verify", lambda **kwargs: (lengths, mask)) calls = [] - linear = object.__new__(GDNState) - linear.commit_accepted_state = lambda **kwargs: calls.append(kwargs) - other = SimpleNamespace(commit_bounded_att_state=lambda indices, steps: calls.append((indices, steps))) - resources = SimpleNamespace(get_bounded_attention_states=lambda role: (linear, other) if role == "main" else ()) + resources = SimpleNamespace(commit_bounded_att_state=lambda ids, steps, **kw: calls.append((ids, steps, kw))) backend = SimpleNamespace( max_draft_step=1, model=SimpleNamespace( @@ -131,12 +129,16 @@ def test_verify_routes_linear_commit_context_and_other_bounded_states(monkeypatc req_manager=SimpleNamespace(req_sampling_params_manager=SimpleNamespace(req_to_next_token_ids=None)), ), ) - batches = ((2, object()),) - utils.verify_mtp_tokens(backend, torch.zeros(4), reqs, starts, positions, linear_state_batches=batches) - assert calls[0]["linear_state_batches"] is batches - assert calls[0]["b_req_idx"] is reqs - assert calls[0]["b_mtp_index"] is positions - assert calls[0]["accepted_index"] is mask - assert calls[0]["verify_width"] == 2 - assert calls[1][0].tolist() == [2, 0] - assert calls[1][1] is lengths + batches = (StateCommitBatch(2),) + utils.verify_mtp_tokens(backend, torch.zeros(4), reqs, starts, positions, state_commit_batches=batches) + assert len(calls) == 1 + ids, steps, kwargs = calls[0] + assert ids.tolist() == [2, 0] + assert steps is lengths + assert kwargs["role"] == "main" + context = kwargs["context"] + assert context.batches is batches + assert context.b_req_idx is reqs + assert context.b_mtp_index is positions + assert context.accepted_index is mask + assert context.verify_width == 2 diff --git a/unit_tests/common/basemodel/attention/linear/test_gdn.py b/unit_tests/common/basemodel/attention/linear/test_gdn.py index a775cca4bc..f9b5b39f76 100644 --- a/unit_tests/common/basemodel/attention/linear/test_gdn.py +++ b/unit_tests/common/basemodel/attention/linear/test_gdn.py @@ -10,6 +10,7 @@ from lightllm.common.attention_state.bounded_att import ConvSSMStateConfig, GDNState from lightllm.common.basemodel.infer_struct import InferStateInfo import lightllm.common.basemodel.attention.linear.gdn as gdn +import lightllm.common.basemodel.attention.linear.state as linear_state import lightllm.common.basemodel.attention.create_linear_utils as linear_create_utils import lightllm.common.basemodel.triton_kernel.linear_att.fla.ops as fla_ops from lightllm.common.basemodel.attention.linear.flashqla import FlashQlaLinearAttBackend @@ -58,6 +59,7 @@ def test_prefill_casts_final_state_to_cache_dtype(monkeypatch, cache_dtype): torch.zeros((1, 1), dtype=cache_dtype), state.infer_state, layer_weight, + (state.b_conv_buffer_idx, state.b_ssm_buffer_idx), ) assert ssm_states.dtype == cache_dtype @@ -69,11 +71,12 @@ def test_decode_state_initializes_normal_layout(): state, infer_state = _create_decode_state(draft_step=0, dynamic_layout=False, b_req_idx=b_req_idx) state.init_state() + forward = state.get_forward(3) - assert state.b_conv_buffer_idx is b_req_idx - assert state.b_ssm_buffer_idx is b_req_idx - assert state.b1_mtp_cu_q_seq_len is None - assert state.b_num_accepted_tokens is None + assert forward.req_indices is b_req_idx + assert forward.state_indices is b_req_idx + assert forward.cu_seqlens is None + assert forward.accepted_token_counts is None @pytest.mark.parametrize("role", ["main", "draft0"]) @@ -89,12 +92,13 @@ def test_decode_state_initializes_fixed_mtp_layout(role: str): ) state.init_state() + forward = state.get_forward(3) - torch.testing.assert_close(state.b1_mtp_cu_q_seq_len, torch.tensor([0, 3, 6], dtype=torch.int32)) - torch.testing.assert_close(state.b_conv_buffer_idx, torch.tensor([2, 4], dtype=torch.int32)) - torch.testing.assert_close(state.b_num_accepted_tokens, torch.tensor([2, 3], dtype=torch.int32)) + torch.testing.assert_close(forward.cu_seqlens, torch.tensor([0, 3, 6], dtype=torch.int32)) + torch.testing.assert_close(forward.req_indices, torch.tensor([2, 4], dtype=torch.int32)) + torch.testing.assert_close(forward.accepted_token_counts, torch.tensor([2, 3], dtype=torch.int32)) torch.testing.assert_close( - state.b_ssm_buffer_idx, + forward.state_indices, torch.tensor([[6, 7, 8], [12, 13, 14]], dtype=torch.int32), ) @@ -112,7 +116,7 @@ def build_params(**kwargs): build_calls.append(kwargs) return expected_cu_q_seq_len, expected_conv_buffer_idx - monkeypatch.setattr(gdn, "build_dynamic_mtp_linear_att_state_params", build_params) + monkeypatch.setattr(linear_state, "build_dynamic_mtp_linear_att_state_params", build_params) state, infer_state = _create_decode_state( draft_step=2, dynamic_layout=True, @@ -122,12 +126,13 @@ def build_params(**kwargs): ) state.init_state() + forward = state.get_forward(3) - assert state.b1_mtp_cu_q_seq_len is expected_cu_q_seq_len - assert state.b_conv_buffer_idx is expected_conv_buffer_idx - torch.testing.assert_close(state.b_num_accepted_tokens, expected_num_accepted_tokens) + assert forward.cu_seqlens is expected_cu_q_seq_len + assert forward.req_indices is expected_conv_buffer_idx + torch.testing.assert_close(forward.accepted_token_counts, expected_num_accepted_tokens) torch.testing.assert_close( - state.b_ssm_buffer_idx, + forward.state_indices, torch.tensor( [[6, 7, 8], [12, 13, 14], [15, 16, 17], [15, 16, 17], [15, 16, 17]], dtype=torch.int32, diff --git a/unit_tests/common/basemodel/test_cuda_graph_autotune.py b/unit_tests/common/basemodel/test_cuda_graph_autotune.py index 71ff5368c8..ee76288d8d 100644 --- a/unit_tests/common/basemodel/test_cuda_graph_autotune.py +++ b/unit_tests/common/basemodel/test_cuda_graph_autotune.py @@ -141,24 +141,30 @@ def test_linear_config_is_local_to_capture_state(capture_env, overlap, kda): from lightllm.common.basemodel.batch_objs import PostLayerOutput from lightllm.common.basemodel.hidden_collector import NoopHiddenCollector from lightllm.common.basemodel.infer_struct import InferStateInfo - from lightllm.common.basemodel.triton_kernel.linear_att.compact_ssm import CompactSSMBackend - from lightllm.common.linear_state import SSMRunConfig + from lightllm.common.attention_state.bounded_att.ssm.compact import CompactSSMState + from lightllm.common.ssm_types import SSMRunConfig from lightllm.common.attention_state.bounded_att import GDNState, KDAState capture_env.args.enable_decode_microbatch_overlap = overlap graph = CudaGraph(1, 2, 1, max_batch_size=2) - manager = object.__new__(KDAState if kda else GDNState) - manager.ssm_backend = object.__new__(CompactSSMBackend) + from unit_tests.common.test_linear_state import make_state + from lightllm.common.attention_state.commit import StateCommit + + manager, resources, _ = make_state(KDAState if kda else GDNState, mode="compact", device="cpu") att_class = KDALinearAttDecodeAttState if kda else GDNLinearAttDecodeAttState states = [] configs = [SSMRunConfig(8, 1), SSMRunConfig(32, 2)] for _ in range(2 if overlap else 1): state = InferStateInfo() state.input_ids = torch.ones(2, dtype=torch.int64) - state.att_state_manager = SimpleNamespace(get_bounded_attention_states=lambda role: (manager,)) + state.att_state_manager = resources state.hidden_collector = NoopHiddenCollector() state.decode_att_state = TritonDecodeAttState(_infer_state_weak_ref=weakref.ref(state)) state.decode_att_state1 = att_class(_infer_state_weak_ref=weakref.ref(state)) + state.decode_att_state1.forwards[manager.state_name] = manager.prepare_decode( + torch.tensor([0, 1], dtype=torch.int32), + torch.tensor([0, 3, 6], dtype=torch.int32), + ) states.append(state) def decode_func(*infer_states): @@ -166,9 +172,10 @@ def decode_func(*infer_states): for index, state in enumerate(infer_states): att = state.decode_att_state1 assert att.infer_state is state - assert att.run_config is None + forward = att.get_forward(0) + assert forward.run_config is None assert (att is states[index].decode_att_state1) == capture_env.capturing - att.run_config = configs[index] if capture_env.capturing else SSMRunConfig(128, 4) + forward.run_config = configs[index] if capture_env.capturing else SSMRunConfig(128, 4) outputs.append(TpPartBaseModel._create_model_output(None, PostLayerOutput(torch.ones(2, 1)), state)) return tuple(outputs) if overlap else outputs[0] @@ -176,13 +183,19 @@ def decode_func(*infer_states): outputs = result if overlap else (result,) runtime_states = [cuda_graph_module._copy_warmup_state(state) for state in states] for state in runtime_states: - state.decode_att_state1.run_config = None + state.decode_att_state1.forwards[manager.state_name] = manager.prepare_decode( + torch.tensor([1, 0], dtype=torch.int32), + torch.tensor([0, 1, 4], dtype=torch.int32), + ) replay = graph.replay(*runtime_states) assert replay == result for index, (state, output) in enumerate(zip(states, outputs)): - assert state.decode_att_state1.run_config is configs[index] - state.decode_att_state1.run_config = SSMRunConfig(64, 4) - assert output.linear_state_commit == manager.create_commit(configs[index]) + forward = state.decode_att_state1.get_forward(0) + assert forward.run_config is configs[index] + assert forward.req_indices.tolist() == [1, 0] + assert forward.cu_seqlens.tolist() == [0, 1, 4] + forward.run_config = SSMRunConfig(64, 4) + assert output.state_commits == (StateCommit(manager.state_name, configs[index]),) unpadded = TpPartBaseModel._create_unpad_decode_model_output(None, output, 1) assert unpadded.logits.shape[0] == 1 - assert unpadded.linear_state_commit is output.linear_state_commit + assert unpadded.state_commits is output.state_commits diff --git a/unit_tests/common/basemodel/triton_kernel/linear_att/test_acceptance_fusion.py b/unit_tests/common/basemodel/triton_kernel/linear_att/test_acceptance_fusion.py index 13ee79dcfb..4694b9fbe6 100644 --- a/unit_tests/common/basemodel/triton_kernel/linear_att/test_acceptance_fusion.py +++ b/unit_tests/common/basemodel/triton_kernel/linear_att/test_acceptance_fusion.py @@ -3,8 +3,10 @@ import triton from lightllm.common.basemodel.triton_kernel.mtp_utils import linear_att_mtp_state_index_update as update -from lightllm.common.basemodel.triton_kernel.linear_att.replayssm import ReplaySSMBackend -from lightllm.common.linear_state import LinearStateManager +from lightllm.common.attention_state.bounded_att.ssm.replay import ReplaySSMState +from types import SimpleNamespace +from lightllm.common.attention_state.commit import MtpCommitContext +from lightllm.common.attention_state.bounded_att.ssm.native import NativeSSMState def make(n, width): @@ -20,24 +22,25 @@ def make(n, width): mtps.extend(range(l)) tensor = lambda x: torch.tensor(x, device="cuda", dtype=torch.int32) args = [tensor(starts), tensor(reqs), tensor(mtps), tensor([1] * len(reqs)), width] - ssm_backend = ReplaySSMBackend.__new__(ReplaySSMBackend) + ssm_backend = ReplaySSMState.__new__(ReplaySSMState) ssm_backend.cursors = tensor([7] * (hold + 1)) ssm_backend.hold_req_id = hold ssm_backend.capacity = max(4, triton.next_power_of_2(width)) ssm_backend.verify_width = width - manager = object.__new__(LinearStateManager) + manager = SimpleNamespace() manager.mtp_state_index = tensor([-9] * (hold + 1)) - manager.ssm_backend = ssm_backend + manager._ssm = ssm_backend return manager, args, (ids, lengths, starts) def accept_unfused(m, args): update(m.mtp_state_index, *args) - m.ssm_backend.accept_updates(args[1][args[0].long()], m.mtp_state_index) + m._ssm.accept_updates(args[1][args[0].long()], m.mtp_state_index) def accept_fused(m, args): - m.commit_accepted_state(*args) + reqs = args[1][args[0].long()] + m._ssm.commit(reqs, None, m.mtp_state_index, MtpCommitContext(*args), "main") @pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") @@ -58,15 +61,15 @@ def test_acceptance_fusion_matches_original_and_cpu_under_graph(): expected_cursor = [7] * len(expected_index) m.mtp_state_index.fill_(-9) ref.mtp_state_index.fill_(-9) - m.ssm_backend.cursors.fill_(7) - ref.ssm_backend.cursors.fill_(7) + m._ssm.cursors.fill_(7) + ref._ssm.cursors.fill_(7) for step in range(24): mask = [] for i, (r, l) in enumerate(zip(ids, lengths)): count = (step + i) % (l + 1) mask.extend([1] * count + [0] * (l - count)) expected_index[r] = count - 1 - if r != m.ssm_backend.hold_req_id: + if r != m._ssm.hold_req_id: expected_cursor[r] += count if width > 1 else 1 args[3].copy_(torch.tensor(mask, device="cuda", dtype=torch.int32)) accept_unfused(ref, args) @@ -75,26 +78,15 @@ def test_acceptance_fusion_matches_original_and_cpu_under_graph(): else: accept_fused(m, args) assert torch.equal(m.mtp_state_index, ref.mtp_state_index) - assert torch.equal(m.ssm_backend.cursors, ref.ssm_backend.cursors) + assert torch.equal(m._ssm.cursors, ref._ssm.cursors) assert m.mtp_state_index.tolist() == expected_index - assert m.ssm_backend.cursors.tolist() == expected_cursor + assert m._ssm.cursors.tolist() == expected_cursor checks += 1 - # Native optional-pointer path and manager dispatch leave metadata-only backend absent. - m.ssm_backend = None + + m._ssm = NativeSSMState(torch.empty(0), width) accept_fused(m, args) assert m.mtp_state_index.tolist() == expected_index - # Non-Replay backends must still receive the accepted request IDs and indexes. - class CompactProbe: - def accept_updates(self, reqs, index): - self.seen = (reqs.clone(), index.clone()) - - m, args, (ids, _, _) = make(7, 4) - probe = CompactProbe() - m.ssm_backend = probe - accept_fused(m, args) - assert probe.seen[0].tolist() == ids - assert torch.equal(probe.seen[1], m.mtp_state_index) print("ACCEPT_EQUIVALENCE_CHECKS=" + str(checks), flush=True) @@ -105,7 +97,7 @@ def test_fused_acceptance_preserves_folds_snapshots_and_slot_reuse(dtype, comput torch.manual_seed(93) initial = torch.randn(2, 8, 4, 32, 32, device="cuda").to(dtype) * 0.01 ssm_backends = [ - ReplaySSMBackend( + ReplaySSMState( initial.clone(), 8, 4, @@ -118,8 +110,8 @@ def test_fused_acceptance_preserves_folds_snapshots_and_slot_reuse(dtype, comput ] managers = [] for ssm_backend in ssm_backends: - manager = object.__new__(LinearStateManager) - manager.ssm_backend = ssm_backend + manager = SimpleNamespace() + manager._ssm = ssm_backend manager.mtp_state_index = torch.zeros(8, dtype=torch.int32, device="cuda") managers.append(manager) reqs = torch.full((4,), 7, dtype=torch.int32, device="cuda") @@ -188,8 +180,8 @@ def cycle(i): assert torch.equal(ssm_backends[0].state, ssm_backends[1].state) if iteration == 19: for manager in managers: - manager.ssm_backend.clear_history(1) - manager.ssm_backend.state[:, 1].zero_() + manager._ssm.clear_history(1) + manager._ssm.state[:, 1].zero_() manager.mtp_state_index[1] = 0 for ssm_backend in ssm_backends: ssm_backend.merge_accepted_updates(reqs[:3]) diff --git a/unit_tests/common/basemodel/triton_kernel/linear_att/test_mtp_state_params.py b/unit_tests/common/basemodel/triton_kernel/linear_att/test_mtp_state_params.py index 32cba5c937..c4c35a78cf 100644 --- a/unit_tests/common/basemodel/triton_kernel/linear_att/test_mtp_state_params.py +++ b/unit_tests/common/basemodel/triton_kernel/linear_att/test_mtp_state_params.py @@ -1,7 +1,7 @@ import pytest import torch -from lightllm.common.linear_state import LinearStateManager +from lightllm.common.attention_state.bounded_att import GDNState, ConvSSMStateConfig from lightllm.common.basemodel.triton_kernel.linear_att.mtp_state_params import ( build_dynamic_mtp_linear_att_state_params, @@ -14,16 +14,24 @@ def test_dynamic_sequences_stay_within_request_capacity_under_graph(request_capa hold = request_capacity reqs = torch.full((token_bucket,), hold, dtype=torch.int32, device="cuda") mtp_index = torch.zeros_like(reqs) - manager = object.__new__(LinearStateManager) - manager.mtp_state_index = torch.arange(request_capacity + 1, dtype=torch.int32, device="cuda") % 4 - manager.ssm_slots_per_req = 4 - manager.ssm_backend = None + config = ConvSSMStateConfig( + logical_layer_ids=(0,), + global_k_heads=1, + global_v_heads=1, + key_dim=2, + value_dim=2, + conv_width=3, + conv_dtype=torch.float32, + ssm_dtype=torch.float32, + ) + manager = GDNState(config, request_capacity + 1, mtp_step=3) + manager.mtp_state_index.copy_(torch.arange(request_capacity + 1, dtype=torch.int32, device="cuda") % 4) sequences = min(request_capacity, token_bucket) def build(): cu, ids = build_dynamic_mtp_linear_att_state_params(reqs, mtp_index, hold) - ids, ssm, counts = manager.prepare_mtp_decode(ids, cu) - return cu, ids, ssm, counts + forward = manager.prepare_decode(ids, cu) + return cu, forward.req_indices, forward.state_indices, forward.accepted_token_counts build() torch.cuda.synchronize() 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 index 165b86b131..4713d1015a 100644 --- a/unit_tests/common/basemodel/triton_kernel/linear_att/test_replayssm.py +++ b/unit_tests/common/basemodel/triton_kernel/linear_att/test_replayssm.py @@ -1,7 +1,7 @@ import pytest import torch -from lightllm.common.basemodel.triton_kernel.linear_att.replayssm import ReplaySSMBackend +from lightllm.common.attention_state.bounded_att.ssm.replay import ReplaySSMState from lightllm.common.basemodel.triton_kernel.linear_att.ssm_autotune import get_configs pytestmark = pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") @@ -30,7 +30,7 @@ def reference(q, k, v, a, b, a_log, bias, state, kda=False, gate_lower_bound=-5. @pytest.mark.parametrize("kda", [False, True], ids=["gdn", "kda"]) def test_fp16_correction_history_preserves_finite_recurrence(capacity, bv, state_dtype, kda): state = torch.zeros(1, 2, 1, 128, 128, dtype=state_dtype, device="cuda") - ssm_backend = ReplaySSMBackend( + ssm_backend = ReplaySSMState( state, capacity, 2, @@ -103,7 +103,7 @@ def test_replay_acceptance_flush_and_graph( h, hv = head_geometry slots = 5 state = (torch.randn(2, slots, hv, kdim, vdim, device="cuda") * 0.01).to(dtype) - ssm_backend = ReplaySSMBackend( + ssm_backend = ReplaySSMState( state, capacity, width, @@ -265,7 +265,7 @@ def test_precomputed_projections_preserve_precision(warps, decay, state_dtype, d state = (torch.randn((1, 3, 4, 128, 128), device="cuda") * 0.1).to(state_dtype) config = {"BV": bv, "num_warps": warps} ssm_backends = [ - ReplaySSMBackend( + ReplaySSMState( state.clone(), 16 if width == 1 else 8, width, @@ -312,7 +312,7 @@ def test_bf16_decode_fold_uses_rounded_checkpoint(): # Automatic fold and explicit materialization must start from the same BF16 # checkpoint, including the token computed inside the automatic-fold kernel. state = torch.full((1, 2, 1, 32, 64), 0.123, device="cuda", dtype=torch.bfloat16) - ssm_backends = [ReplaySSMBackend(state.clone(), 16, 1, torch.float32) for _ in range(2)] + ssm_backends = [ReplaySSMState(state.clone(), 16, 1, torch.float32) for _ in range(2)] reqs = torch.tensor([0], device="cuda", dtype=torch.int32) for ssm_backend in ssm_backends: ssm_backend.raw_keys.zero_() @@ -343,7 +343,7 @@ def test_bf16_decode_fold_uses_rounded_checkpoint(): def test_dynamic_verify_lengths_fill_history_before_folding(dtype, compute_mode, capture_metadata): torch.manual_seed(31) state = (torch.randn(2, 4, 4, 128, 128, device="cuda") * 0.01).to(dtype) - ssm_backend = ReplaySSMBackend( + ssm_backend = ReplaySSMState( state, 8, 4, dtype, num_key_heads=2, compute_mode=compute_mode, run_config={"BV": 32, "num_warps": 1} ) expected = state.float().clone() @@ -431,7 +431,7 @@ def step(): def test_snapshot_without_history_copies_active_state(): state = torch.randn((2, 3, 2, 32, 64), device="cuda", dtype=torch.bfloat16) - ssm_backend = ReplaySSMBackend(state, 4, 3) + ssm_backend = ReplaySSMState(state, 4, 3) snapshot = ssm_backend.get_accepted_state(1) torch.testing.assert_close(snapshot, state[:, 1], rtol=0, atol=0) assert ssm_backend.cursors[1].item() == 0 @@ -439,12 +439,12 @@ def test_snapshot_without_history_copies_active_state(): @pytest.mark.parametrize("dtype", [torch.float32, torch.bfloat16]) def test_compact_preserves_accepted_prefix(dtype, run_config=None, kda=False): - from lightllm.common.basemodel.triton_kernel.linear_att.compact_ssm import CompactSSMBackend + from lightllm.common.attention_state.bounded_att.ssm.compact import CompactSSMState torch.manual_seed(17) hv, kd, vd, width = 4, 128 if kda else 64, 128 if kda else 64, 3 state = torch.randn(2, 4, hv, kd, vd, device="cuda", dtype=dtype) * 0.01 - ssm_backend = CompactSSMBackend(state, width, torch.bfloat16, num_key_heads=hv, kda=kda, run_config=run_config) + ssm_backend = CompactSSMState(state, width, torch.bfloat16, num_key_heads=hv, kda=kda, run_config=run_config) 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) @@ -479,7 +479,7 @@ def test_compact_preserves_accepted_prefix(dtype, run_config=None, kda=False): @pytest.mark.parametrize("dtype", [torch.float32, torch.bfloat16]) @pytest.mark.parametrize("width", [3, 4]) def test_compact_matches_native_mtp(dtype, width, run_config=None): - from lightllm.common.basemodel.triton_kernel.linear_att.compact_ssm import CompactSSMBackend + from lightllm.common.attention_state.bounded_att.ssm.compact import CompactSSMState from lightllm.common.basemodel.triton_kernel.linear_att.mtp_fused_recurrent import ( mtp_fused_recurrent_gated_delta_rule, ) @@ -500,7 +500,7 @@ def test_compact_matches_native_mtp(dtype, width, run_config=None): 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) - ssm_backend = CompactSSMBackend(state, width, torch.bfloat16, num_key_heads=h, run_config=run_config) + ssm_backend = CompactSSMState(state, width, torch.bfloat16, num_key_heads=h, run_config=run_config) for _ in range(32): q.normal_() k.normal_() @@ -570,7 +570,7 @@ def test_checkpoint_restores_pending_history_and_accepted_conv(save_big, dtype): state.mtp_state_index[1] = 1 state.conv_state.normal_() state.ssm_state.zero_() - backend = state.ssm_backend + backend = state._ssm backend.raw_keys.normal_() backend.raw_values.normal_() backend.gates.uniform_(-0.1, -0.01) diff --git a/unit_tests/common/basemodel/triton_kernel/linear_att/test_replayssm_execution.py b/unit_tests/common/basemodel/triton_kernel/linear_att/test_replayssm_execution.py index e172b6495c..6da95389ce 100644 --- a/unit_tests/common/basemodel/triton_kernel/linear_att/test_replayssm_execution.py +++ b/unit_tests/common/basemodel/triton_kernel/linear_att/test_replayssm_execution.py @@ -1,11 +1,12 @@ import pytest import torch -from lightllm.common.basemodel.triton_kernel.linear_att.replayssm import ReplaySSMBackend +from lightllm.common.attention_state.bounded_att.ssm.replay import ReplaySSMState from lightllm.common.basemodel.triton_kernel.linear_att.mtp_state_params import ( build_dynamic_mtp_linear_att_state_params, ) -from lightllm.common.linear_state import LinearStateManager +from types import SimpleNamespace +from lightllm.common.attention_state.commit import MtpCommitContext pytestmark = pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") @@ -29,10 +30,10 @@ def test_request_manager_accept_with_replay_graph(dtype, capture_metadata, varle compute_mode="precompute", run_config=config, ) - for cls in [ReplaySSMBackend, ReplaySSMBackend] + for cls in [ReplaySSMState, ReplaySSMState] ] - manager = object.__new__(LinearStateManager) - manager.ssm_backend = ssm_backends[1] + manager = SimpleNamespace() + manager._ssm = ssm_backends[1] manager.mtp_state_index = torch.zeros(5, device="cuda", dtype=torch.int32) rows = torch.full((12,), 4, device="cuda", dtype=torch.int32) mtp = torch.zeros_like(rows) @@ -80,7 +81,13 @@ def forward_accept(index, cu, reqs): if index == 0: ssm_backend.accept_updates(reqs, accepted) else: - manager.commit_accepted_state(starts, rows, mtp, flags, 4) + manager._ssm.commit( + rows[starts.long()], + None, + manager.mtp_state_index, + MtpCommitContext(starts, rows, mtp, flags, 4), + "main", + ) return out inputs(0) @@ -142,7 +149,7 @@ def test_decode_accept_order_preserves_forward_snapshot(dtype, capture_metadata, torch.manual_seed(731) initial = torch.randn(3, 3, 4, 128, 128, device="cuda", dtype=dtype) * 0.01 ssm_backends = [ - ReplaySSMBackend( + ReplaySSMState( initial.clone(), 8, 1, @@ -237,7 +244,7 @@ def test_cross_layer_fold_with_state_offset_above_int32(): sparse.copy_(dense) config = {"BV": 32, "num_warps": 1, "num_stages": 1} ssm_backends = [ - ReplaySSMBackend(state, 8, 4, torch.bfloat16, num_key_heads=4, run_config=config) for state in [dense, sparse] + ReplaySSMState(state, 8, 4, torch.bfloat16, num_key_heads=4, run_config=config) for state in [dense, sparse] ] q = torch.randn(1, 8, 4, 128, device="cuda", dtype=torch.bfloat16) k = torch.randn_like(q) @@ -280,7 +287,7 @@ def test_small_batch_matches_larger_graph_padding(batch, dtype, kda): initial = torch.randn(2, 14, 4, 128, 128, device="cuda", dtype=dtype) * 0.01 config = {"BV": 32, "num_warps": 1, "num_stages": 1} ssm_backends = [ - ReplaySSMBackend(initial.clone(), 8, 4, num_key_heads=2, kda=kda, compute_mode="precompute", run_config=config) + ReplaySSMState(initial.clone(), 8, 4, num_key_heads=2, kda=kda, compute_mode="precompute", run_config=config) for _ in range(2) ] reqs = torch.arange(batch, device="cuda", dtype=torch.int32) diff --git a/unit_tests/common/basemodel/triton_kernel/linear_att/test_replayssm_grouped_keys.py b/unit_tests/common/basemodel/triton_kernel/linear_att/test_replayssm_grouped_keys.py index 7546fd38f9..78f80de800 100644 --- a/unit_tests/common/basemodel/triton_kernel/linear_att/test_replayssm_grouped_keys.py +++ b/unit_tests/common/basemodel/triton_kernel/linear_att/test_replayssm_grouped_keys.py @@ -3,8 +3,8 @@ import pytest import torch -from lightllm.common.basemodel.triton_kernel.linear_att.replayssm import ReplaySSMBackend -from lightllm.common.basemodel.triton_kernel.linear_att.compact_ssm import CompactSSMBackend +from lightllm.common.attention_state.bounded_att.ssm.replay import ReplaySSMState +from lightllm.common.attention_state.bounded_att.ssm.compact import CompactSSMState pytestmark = pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") @@ -20,8 +20,8 @@ def test_grouped_keys_graph_fold_partial_accept(mode, group, dtype): def make_ssm_backend(key_heads): if mode == "replay": - return ReplaySSMBackend(state.clone(), 16, width, num_key_heads=key_heads) - return CompactSSMBackend(state.clone(), width, torch.bfloat16, num_key_heads=key_heads) + return ReplaySSMState(state.clone(), 16, width, num_key_heads=key_heads) + return CompactSSMState(state.clone(), width, torch.bfloat16, num_key_heads=key_heads) ssm_backends = [make_ssm_backend(hv), make_ssm_backend(h)] assert ssm_backends[0].keys.numel() == ssm_backends[1].keys.numel() * group diff --git a/unit_tests/common/basemodel/triton_kernel/linear_att/test_replayssm_parallel.py b/unit_tests/common/basemodel/triton_kernel/linear_att/test_replayssm_parallel.py index 7d869ac2c4..25bcb15126 100644 --- a/unit_tests/common/basemodel/triton_kernel/linear_att/test_replayssm_parallel.py +++ b/unit_tests/common/basemodel/triton_kernel/linear_att/test_replayssm_parallel.py @@ -1,7 +1,7 @@ import pytest import torch -from lightllm.common.basemodel.triton_kernel.linear_att.replayssm import ReplaySSMBackend +from lightllm.common.attention_state.bounded_att.ssm.replay import ReplaySSMState pytestmark = pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") @@ -11,7 +11,7 @@ def test_derived_fold_commits_only_accepted_history(capacity): torch.manual_seed(731) state = torch.randn(2, 5, 4, 128, 128, device="cuda", dtype=torch.bfloat16) * 0.02 initial = state.clone() - ssm_backend = ReplaySSMBackend(state, capacity, 5, num_key_heads=2, compute_mode="parallel", fold_method="derived") + ssm_backend = ReplaySSMState(state, capacity, 5, num_key_heads=2, compute_mode="parallel", fold_method="derived") ssm_backend.keys.fill_(float("nan")) ssm_backend.deltas.fill_(float("nan")) ssm_backend.gates.fill_(float("nan")) @@ -46,13 +46,13 @@ def test_derived_fold_commits_only_accepted_history(capacity): def test_parallel_rejects_unsupported_shapes(fold): state = torch.zeros(1, 3, 4, 32, 64, device="cuda", dtype=torch.bfloat16) with pytest.raises(AssertionError, match="128x128"): - ReplaySSMBackend(state, 16, 5, num_key_heads=2, compute_mode="parallel", fold_method=fold) + ReplaySSMState(state, 16, 5, num_key_heads=2, compute_mode="parallel", fold_method=fold) def test_derived_fold_requires_parallel_solver(): state = torch.zeros(1, 3, 4, 128, 128, device="cuda", dtype=torch.bfloat16) with pytest.raises(AssertionError): - ReplaySSMBackend(state, 16, 5, num_key_heads=2, fold_method="derived") + ReplaySSMState(state, 16, 5, num_key_heads=2, fold_method="derived") @pytest.mark.parametrize("width", [2, 3, 5, 8, 16]) @@ -79,7 +79,7 @@ def test_parallel_derived_continuous_state(capacity, width, capture): torch.manual_seed(731) state = torch.randn(2, 4, 4, 128, 128, device="cuda", dtype=torch.bfloat16) * 0.02 initial = state.clone() - ssm_backend = ReplaySSMBackend( + ssm_backend = ReplaySSMState( state, capacity, width, num_key_heads=2, compute_mode="parallel", fold_method="derived" ) for history in ( diff --git a/unit_tests/common/basemodel/triton_kernel/linear_att/test_ssm_autotune.py b/unit_tests/common/basemodel/triton_kernel/linear_att/test_ssm_autotune.py index 03c13144f6..22862ea29f 100644 --- a/unit_tests/common/basemodel/triton_kernel/linear_att/test_ssm_autotune.py +++ b/unit_tests/common/basemodel/triton_kernel/linear_att/test_ssm_autotune.py @@ -6,9 +6,10 @@ import torch from lightllm.common.basemodel.triton_kernel.linear_att import ssm_autotune as tuning -from lightllm.common.basemodel.triton_kernel.linear_att.replayssm import ReplaySSMBackend -from lightllm.common.basemodel.triton_kernel.linear_att.compact_ssm import CompactSSMBackend -from lightllm.common.linear_state import LinearStateManager +from lightllm.common.attention_state.bounded_att.ssm.replay import ReplaySSMState +from lightllm.common.attention_state.bounded_att.ssm.compact import CompactSSMState +from lightllm.common.attention_state.commit import MtpCommitContext, StateCommitBatch +from lightllm.common.attention_state.bounded_att.ssm.forward import SSMForward from lightllm.common.kernel_config import KernelConfigs from lightllm.common import kernel_config as kernel_config_module from lightllm.common.triton_utils import autotuner as autotuner_module @@ -19,9 +20,9 @@ def make_inputs(mode, dtype=torch.bfloat16, device="cpu", width=4, compute_mode= layers, slots, batch, h, hv, kd, vd = 2, 9, 3, 2, 4, 32, 64 state = torch.randn(layers, slots, hv, kd, vd, dtype=dtype, device=device) * 0.01 if mode == "replay": - ssm_backend = ReplaySSMBackend(state, 8, width, num_key_heads=h, compute_mode=compute_mode, kda=kda) + ssm_backend = ReplaySSMState(state, 8, width, num_key_heads=h, compute_mode=compute_mode, kda=kda) else: - ssm_backend = CompactSSMBackend(state, width, torch.bfloat16, num_key_heads=h, kda=kda) + ssm_backend = CompactSSMState(state, width, torch.bfloat16, num_key_heads=h, kda=kda) tokens = batch * width packed = torch.randn(1, tokens, 2 * h * kd + hv * vd, dtype=torch.bfloat16, device=device) q, k, v = [ @@ -334,7 +335,8 @@ def test_two_verify_graphs_accept_with_their_own_configs(monkeypatch, mode, proj from lightllm.common.attention_state.bounded_att import KDAState, GDNState manager = object.__new__(KDAState if kda else GDNState) - manager.ssm_backend = ssm_backend + manager._ssm = ssm_backend + manager.state_name = "main_kda" if kda else "main_gdn" manager.mtp_state_index = torch.zeros(9, dtype=torch.int32, device="cuda") buckets = [] for index, ids in enumerate(([0], [1, 3])): @@ -374,14 +376,17 @@ def forward(target): state.att_state_manager = SimpleNamespace(get_bounded_attention_states=lambda role: (manager,)) state.hidden_collector = SimpleNamespace(finish_output=lambda **kw: ModelMtpOutputCollector()) att_class = KDALinearAttDecodeAttState if kda else GDNLinearAttDecodeAttState - state.decode_att_state1 = att_class(_infer_state_weak_ref=weakref.ref(state), run_config=config) + state.decode_att_state1 = att_class( + _infer_state_weak_ref=weakref.ref(state), + forwards={manager.state_name: SSMForward(manager, reqs, reqs, cu, None, config)}, + ) output = TpPartBaseModel._create_model_output( None, PostLayerOutput(logits=outputs[0].view(count * 4, -1)), state ) - commit = output.linear_state_commit - assert (commit is not None) == (mode == "compact") + commits = output.state_commits + assert bool(commits) == (mode == "compact") output = TpPartBaseModel._create_unpad_decode_model_output(None, output, len(ids) * 4) - assert output.linear_state_commit is commit + assert output.state_commits is commits buckets.append((graph, outputs, output, ref, args, reqs, cu, ids)) # Published contexts and captured launches must survive later changes to # the tuner's dictionaries, including when two graphs replay alternately. @@ -405,14 +410,10 @@ def forward(target): torch.testing.assert_close(outputs[layer], expected, atol=0, rtol=0) counts = [1 + (iteration + i) % 4 for i in range(3)] flags = torch.tensor([int(j < n) for n in counts for j in range(4)], dtype=torch.int32, device="cuda") - manager.commit_accepted_state( - starts, - rows, - mtp, - flags, - 4, - linear_state_batches=tuple((len(b[-1]), b[2].linear_state_commit) for b in buckets), + context = MtpCommitContext( + starts, rows, mtp, flags, 4, tuple(StateCommitBatch(len(b[-1]), b[2].state_commits) for b in buckets) ) + manager.commit_bounded_att_state(rows[starts.long()], torch.tensor(counts, device="cuda"), context=context) for _, _, _, ref, _, reqs, _, ids in buckets: ref.accept_updates(reqs[:-1], manager.mtp_state_index) torch.testing.assert_close(ssm_backend.state[:, ids], ref.state[:, ids], atol=0, rtol=0) diff --git a/unit_tests/common/basemodel/triton_kernel/linear_att/test_ssm_autotune_memory.py b/unit_tests/common/basemodel/triton_kernel/linear_att/test_ssm_autotune_memory.py index 398100ce7c..8ba764390e 100644 --- a/unit_tests/common/basemodel/triton_kernel/linear_att/test_ssm_autotune_memory.py +++ b/unit_tests/common/basemodel/triton_kernel/linear_att/test_ssm_autotune_memory.py @@ -2,13 +2,13 @@ import torch from lightllm.common.basemodel.triton_kernel.linear_att import ssm_autotune as tuning -from lightllm.common.basemodel.triton_kernel.linear_att.replayssm import ReplaySSMBackend +from lightllm.common.attention_state.bounded_att.ssm.replay import ReplaySSMState @pytest.mark.parametrize("width", [1, 5]) def test_scratch_budget_is_independent_of_model_layer_count(width): state = torch.empty((48, 257, 48, 128, 128), dtype=torch.bfloat16, device="meta") - ssm_backend = ReplaySSMBackend(state, 16, width, num_key_heads=16) + ssm_backend = ReplaySSMState(state, 16, width, num_key_heads=16) tokens = 256 * width inputs = [ torch.empty(shape, dtype=torch.bfloat16, device="meta") @@ -16,7 +16,7 @@ def test_scratch_budget_is_independent_of_model_layer_count(width): ] required = tuning.scratch_size_bytes(ssm_backend, 256, 256, inputs) assert required < 1024 ** 3 - single_layer = ReplaySSMBackend(state[:1], 16, width, num_key_heads=16) + single_layer = ReplaySSMState(state[:1], 16, width, num_key_heads=16) assert required == tuning.scratch_size_bytes(single_layer, 256, 256, inputs) @@ -27,7 +27,7 @@ def test_tuning_budget_falls_back_without_touching_serving_state(monkeypatch, tm from lightllm.common.triton_utils.autotuner import Autotuner, AutotuneKernelType state = torch.zeros((2, 4, 2, 32, 32), dtype=torch.bfloat16, device="cuda") - ssm_backend = ReplaySSMBackend(state, 8, width, num_key_heads=2) + ssm_backend = ReplaySSMState(state, 8, width, num_key_heads=2) for tensor in vars(ssm_backend).values(): if isinstance(tensor, torch.Tensor): tensor.zero_() diff --git a/unit_tests/common/test_hybrid_req_manager_pd_sharing.py b/unit_tests/common/test_hybrid_req_manager_pd_sharing.py index f634a833e7..cd0006dd45 100644 --- a/unit_tests/common/test_hybrid_req_manager_pd_sharing.py +++ b/unit_tests/common/test_hybrid_req_manager_pd_sharing.py @@ -20,8 +20,8 @@ def _assert_received_manager_roundtrip(payload, state_name, mode, expected): state = getattr(resources, state_name) assert state.big_pages_pool is state.small_pages_pool is None assert resources.full_att_allocator is None - if state.ssm_backend is not None: - assert state.ssm_backend.state is state.ssm_state + if state._ssm is not None: + assert state._ssm.state is state.ssm_state layout = resources.get_pd_operator().get_pd_bounded_att_page_layouts()[state_name] block = torch.empty((1, layout.byte_end), dtype=torch.uint8, device="cuda") state.write_to_pd_page(1, block, 0, 8, layout) @@ -32,8 +32,8 @@ def _assert_received_manager_roundtrip(payload, state_name, mode, expected): assert state.mtp_state_index[2] == 0 assert state.mtp_state_index[1] == 1 if mode == "replay": - assert state.ssm_backend.cursors[2] == 0 - assert state.ssm_backend.cursors[1] == 1 + assert state._ssm.cursors[2] == 0 + assert state._ssm.cursors[1] == 1 torch.cuda.synchronize() @@ -81,7 +81,7 @@ def test_pd_export_copy_keeps_worker_pools_and_runtime_identity(cls, mode): assert copy.big_pages_pool is copy.small_pages_pool is None assert copy.conv_state is state.conv_state assert copy.ssm_state is state.ssm_state - assert copy.ssm_backend is state.ssm_backend + assert copy._ssm is state._ssm assert copy.mtp_state_index is state.mtp_state_index assert all(vars(state)[name] is value for name, value in before.items()) @@ -100,11 +100,11 @@ def test_shared_attention_state_roundtrip_preserves_worker_resources( state.ssm_state[:, slot].fill_(slot + 1) state.mtp_state_index[1], state.mtp_state_index[2] = 1, 2 if mode == "replay": - state.ssm_backend.raw_keys.zero_() - state.ssm_backend.raw_values.zero_() - state.ssm_backend.gates.fill_(-0.5) - state.ssm_backend.betas.zero_() - state.ssm_backend.cursors[1], state.ssm_backend.cursors[2] = 1, 3 + state._ssm.raw_keys.zero_() + state._ssm.raw_values.zero_() + state._ssm.gates.fill_(-0.5) + state._ssm.betas.zero_() + state._ssm.cursors[1], state._ssm.cursors[2] = 1, 3 expected = [tensor.cpu().tolist() for tensor in state.get_accepted_state(1)] before = vars(state).copy() source_conv, source_ssm = state.conv_state[:, 1].clone(), state.ssm_state[:, state.ssm_slots_per_req].clone() diff --git a/unit_tests/common/test_linear_state.py b/unit_tests/common/test_linear_state.py index 55de82d4e1..25f00f42e3 100644 --- a/unit_tests/common/test_linear_state.py +++ b/unit_tests/common/test_linear_state.py @@ -43,11 +43,11 @@ def test_checkpoint_and_pd_export_same_accepted_prefix_and_restore_reused_slot(c expected_conv = state.conv_state[:, src, ..., offset : offset + 3].clone() expected_ssm = state.ssm_state[:, src * state.ssm_slots_per_req + (offset if mode == "native" else 0)].clone() if mode == "replay": - state.ssm_backend.raw_keys.zero_() - state.ssm_backend.raw_values.zero_() - state.ssm_backend.gates.fill_(-0.5) - state.ssm_backend.betas.zero_() - state.ssm_backend.cursors[src] = 1 + state._ssm.raw_keys.zero_() + state._ssm.raw_values.zero_() + state._ssm.gates.fill_(-0.5) + state._ssm.betas.zero_() + state._ssm.cursors[src] = 1 expected_ssm = (expected_ssm.float() * torch.exp(torch.tensor(-0.5))).to(dtype) before = state.ssm_state.clone() block = torch.empty((1, layout.byte_end), dtype=torch.uint8, device="cuda") @@ -60,7 +60,7 @@ def test_checkpoint_and_pd_export_same_accepted_prefix_and_restore_reused_slot(c torch.testing.assert_close(pool.ssm_state[0], expected_ssm.cpu(), rtol=2e-5, atol=2e-5) assert torch.equal(state.ssm_state, before) if mode == "replay": - assert state.ssm_backend.cursors[src].item() == 1 + assert state._ssm.cursors[src].item() == 1 for restore in ( lambda: state.restore_bounded_att_runtime_from_pages([dst], [8], "small", [0]), @@ -70,7 +70,7 @@ def test_checkpoint_and_pd_export_same_accepted_prefix_and_restore_reused_slot(c if state.mtp_state_index is not None: state.mtp_state_index[dst] = step if mode == "replay": - state.ssm_backend.cursors[dst] = 3 + state._ssm.cursors[dst] = 3 restore() actual = state.get_accepted_state(dst) torch.testing.assert_close(actual.conv, expected_conv, rtol=0, atol=0) @@ -78,7 +78,7 @@ def test_checkpoint_and_pd_export_same_accepted_prefix_and_restore_reused_slot(c if state.mtp_state_index is not None: assert state.mtp_state_index[dst].item() == 0 if mode == "replay": - assert state.ssm_backend.cursors[dst].item() == 0 + assert state._ssm.cursors[dst].item() == 0 state.reset_bounded_att_state([dst]) assert not state.conv_state[:, dst].any() start = dst * state.ssm_slots_per_req @@ -92,7 +92,11 @@ def copy_on_cpu(*, copies): def test_pd_block_roundtrip_preserves_replicated_heads_and_accepted_mtp_state(monkeypatch): # Exercise layout and accepted-state selection independently of the CUDA copy kernel. - monkeypatch.setattr(import_module(GDNState.__module__), "copy_bounded_att_tensors", copy_on_cpu) + monkeypatch.setattr( + import_module("lightllm.common.attention_state.bounded_att.conv_ssm_state"), + "copy_bounded_att_tensors", + copy_on_cpu, + ) config = ConvSSMStateConfig( logical_layer_ids=(0,), global_k_heads=1, diff --git a/unit_tests/common/test_linear_state_contract.py b/unit_tests/common/test_linear_state_contract.py index 217f27a0cb..391a0aa598 100644 --- a/unit_tests/common/test_linear_state_contract.py +++ b/unit_tests/common/test_linear_state_contract.py @@ -5,7 +5,10 @@ import pytest import torch -from lightllm.common.linear_state import LinearStateManager, SSMRunConfig +from lightllm.common.ssm_types import SSMRunConfig +from lightllm.common.attention_state.commit import MtpCommitContext, StateCommit, StateCommitBatch +from lightllm.common.attention_state.bounded_att import KDAState +from unit_tests.common.test_linear_state import make_state from lightllm.common.attention_state.bounded_att import ConvSSMStateConfig, GDNState from lightllm.models.glm5_next.config import Glm5NextResourceConfig from lightllm.server.core.objs.start_args_type import StartArgs @@ -82,19 +85,18 @@ def test_model_configuration_retains_nondefault_kda_gate_bound(): assert config.gate_lower_bound == -2.5 -@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") @pytest.mark.parametrize("mode", ["compact", "replay"]) def test_kda_backend_receives_gate_bound_at_construction(linear_config, mode): - owner = LinearStateManager(linear_config, 3, mtp_step=2, ssm_state_mode=mode, kda=True, gate_lower_bound=-2.5) - assert owner.ssm_backend.gate_lower_bound == -2.5 + owner = KDAState(linear_config, 3, mtp_step=2, ssm_state_mode=mode, gate_lower_bound=-2.5, device="cpu") + assert owner._ssm.gate_lower_bound == -2.5 @pytest.mark.parametrize("mode", ["compact", "replay", "parallel"]) def test_ssm_backend_publishes_detached_run_config(mode): - from lightllm.common.basemodel.triton_kernel.linear_att.replayssm import ReplaySSMBackend - from lightllm.common.basemodel.triton_kernel.linear_att.compact_ssm import CompactSSMBackend + from lightllm.common.attention_state.bounded_att.ssm.replay import ReplaySSMState + from lightllm.common.attention_state.bounded_att.ssm.compact import CompactSSMState - ssm_backend = object.__new__(CompactSSMBackend if mode == "compact" else ReplaySSMBackend) + ssm_backend = object.__new__(CompactSSMState if mode == "compact" else ReplaySSMState) tuning = {"BV": 32, "num_warps": 2} ssm_backend.run_config = tuning ssm_backend._config_is_fixed = True @@ -108,31 +110,220 @@ def test_ssm_backend_publishes_detached_run_config(mode): config.BV = 128 +def make_commit_context(batches): + return MtpCommitContext( + torch.tensor([0, 4, 8]), + torch.tensor([2] * 4 + [0] * 4 + [1] * 4), + torch.arange(4).repeat(3), + torch.ones(12, dtype=torch.int32), + 4, + batches, + ) + + def test_compact_commits_each_unpadded_partition_with_its_verify_config(monkeypatch): - from lightllm.common.basemodel.triton_kernel import mtp_utils - from lightllm.common.basemodel.triton_kernel.linear_att.compact_ssm import CompactSSMBackend + from lightllm.common.attention_state.bounded_att.ssm import base - owner = object.__new__(LinearStateManager) - owner.ssm_backend = object.__new__(CompactSSMBackend) - owner.mtp_state_index = torch.tensor([0, 1, 2, 0], dtype=torch.int32) + owner, _, _ = make_state(mode="compact", step=3, device="cpu") + reqs = torch.tensor([2, 0, 1], dtype=torch.int32) configs = [SSMRunConfig(8, 1), SSMRunConfig(32, 2)] - commits = [owner.create_commit(config) for config in configs] - owner.ssm_backend.run_config = {"BV": 128, "num_warps": 4} + commits = [] + for subset, config in zip((reqs[:1], reqs[1:]), configs): + forward = owner.prepare_decode(subset, torch.arange(subset.numel() + 1) * 4) + forward.run_config = config + commits.append(forward.finish()) + owner._ssm.run_config = {"BV": 128, "num_warps": 4} calls = [] - monkeypatch.setattr(mtp_utils, "linear_att_mtp_state_index_update", lambda **kwargs: None) + monkeypatch.setattr(base, "linear_att_mtp_state_index_update", lambda **kwargs: None) monkeypatch.setattr( - owner.ssm_backend, + owner._ssm, "accept_updates", lambda reqs, accepted, run_config: calls.append((reqs.tolist(), accepted, run_config)), ) - owner.commit_accepted_state( - torch.tensor([0, 4, 8]), - torch.tensor([2] * 4 + [0] * 4 + [1] * 4), - torch.arange(4).repeat(3), - torch.ones(12, dtype=torch.int32), - 4, - linear_state_batches=((1, commits[0]), (0, None), (2, commits[1])), + context = make_commit_context( + (StateCommitBatch(1, (commits[0],)), StateCommitBatch(0), StateCommitBatch(2, (commits[1],))) ) + owner.commit_bounded_att_state(reqs, torch.tensor([1, 2, 3]), context=context) assert [reqs for reqs, _, _ in calls] == [[2], [0, 1]] assert all(accepted is owner.mtp_state_index for _, accepted, _ in calls) assert [config for _, _, config in calls] == configs + + +@pytest.mark.parametrize("failure", ["missing", "wrong_owner", "mutable_config", "no_context"]) +def test_invalid_compact_receipt_rejected_before_state_changes(monkeypatch, failure): + owner, _, _ = make_state(mode="compact", step=3, device="cpu") + owner.mtp_state_index.fill_(2) + name = "draft0_other_state" if failure == "wrong_owner" else owner.state_name + config = {"BV": 8, "num_warps": 1} if failure == "mutable_config" else SSMRunConfig(8, 1) + commits = () if failure == "missing" else (StateCommit(name, config),) + context = None if failure == "no_context" else make_commit_context((StateCommitBatch(3, commits),)) + with pytest.raises(ValueError, match="compact"): + owner.commit_bounded_att_state(torch.tensor([2, 0, 1]), torch.tensor([1, 2, 3]), context=context) + assert owner.mtp_state_index.tolist() == [2] * 4 + + +def test_commit_partitions_and_resource_identity_are_validated(): + with pytest.raises(ValueError, match="nonnegative"): + StateCommitBatch(-1) + commit = StateCommit("main", SSMRunConfig(8, 1)) + with pytest.raises(ValueError, match="duplicate"): + StateCommitBatch(1, (commit, commit)) + context = make_commit_context((StateCommitBatch(3, (commit,)),)) + with pytest.raises(ValueError, match="unknown state"): + context.validate(3, {"draft0"}) + with pytest.raises(ValueError, match="counts differ"): + context.validate(2, {"main"}) + with pytest.raises(ValueError, match="partition"): + make_commit_context((StateCommitBatch(2),)).validate(3, {"main"}) + + +@pytest.mark.parametrize("cls", [GDNState, KDAState]) +def test_native_accept_restore_and_reuse_preserve_other_requests(cls): + owner, _, _ = make_state(cls, step=2, device="cpu") + owner.conv_state.copy_(torch.arange(5)) + for slot in range(owner.ssm_state.shape[1]): + owner.ssm_state[:, slot].fill_(slot) + before = owner.ssm_state.clone() + owner.commit_bounded_att_state(torch.tensor([1]), torch.tensor([3], dtype=torch.int32)) + accepted = owner.get_accepted_state(1) + assert torch.equal(accepted.conv[0, 0], torch.tensor([2, 3, 4], dtype=accepted.conv.dtype)) + assert torch.all(accepted.ssm == 5) + owner.restore(2, *accepted) + restored = owner.get_accepted_state(2) + assert torch.equal(restored.conv, accepted.conv) + assert torch.equal(restored.ssm, accepted.ssm) + owner.reset_bounded_att_state([2]) + assert torch.equal(owner.ssm_state[:, :6], before[:, :6]) + assert torch.equal(owner.ssm_state[:, 9:], before[:, 9:]) + assert not owner.ssm_state[:, 6:9].any() + assert owner.mtp_state_index.tolist() == [0, 2, 0, 0] + + +@pytest.mark.parametrize("cls", [GDNState, KDAState]) +def test_replay_fused_acceptance_keeps_cursors_inside_mode(monkeypatch, cls): + from lightllm.common.attention_state.bounded_att.ssm import base + + owner, _, _ = make_state(cls, mode="replay", step=3, device="cpu") + calls = [] + monkeypatch.setattr(base, "linear_att_mtp_state_index_update", lambda **kwargs: calls.append(kwargs)) + monkeypatch.setattr(owner._ssm, "accept_updates", lambda *a: pytest.fail("unfused acceptance")) + owner.commit_bounded_att_state(torch.tensor([2, 0, 1]), torch.tensor([1, 2, 3]), context=make_commit_context(())) + assert len(calls) == 1 + assert calls[0]["replay_cursors"] is owner._ssm.cursors + assert calls[0]["req_to_mtp_state_index"] is owner.mtp_state_index + + +@pytest.mark.parametrize("change", ["owner", "shape", "optional", "dtype"]) +def test_graph_layout_mismatch_rejected_before_copy(change): + from copy import copy + + owner, _, _ = make_state(mode="native", device="cpu") + target = owner.prepare_decode(torch.tensor([0, 1], dtype=torch.int32), torch.tensor([0, 3, 6])) + source = owner.prepare_decode(torch.tensor([1, 2], dtype=torch.int32), torch.tensor([0, 2, 5])) + if change == "owner": + source.resource = copy(owner) + elif change == "shape": + source.state_indices = source.state_indices[:1] + elif change == "optional": + source.accepted_token_counts = None + else: + source.cu_seqlens = source.cu_seqlens.int() + with pytest.raises(ValueError, match="resource|layout"): + target.copy_inputs(source) + assert target.req_indices.tolist() == [0, 1] + assert target.cu_seqlens.tolist() == [0, 3, 6] + + +def test_forward_contexts_select_each_layer_resource_and_publish_both_commits(): + import weakref + from dataclasses import replace + from lightllm.common.attention_state import AttentionStateManager + from lightllm.common.basemodel.attention.linear.gdn import GDNLinearAttDecodeAttState + from lightllm.common.basemodel.batch_objs import PostLayerOutput + from lightllm.common.basemodel.basemodel import TpPartBaseModel + from lightllm.common.basemodel.hidden_collector import NoopHiddenCollector + from lightllm.common.basemodel.infer_struct import InferStateInfo + + first, _, _ = make_state(mode="compact", device="cpu") + second = GDNState( + replace(first.config, logical_layer_ids=(3,), ssm_dtype=torch.bfloat16), + 4, + mtp_step=2, + device="cpu", + ssm_state_mode="compact", + ) + resources = AttentionStateManager(request_capacity=4) + resources.register_state("main", first) + resources.register_state("main", second) + resources.register_state("draft0", GDNState(first.config, 4, mtp_step=2, device="cpu")) + resources.finish_state_registration() + infer = InferStateInfo() + infer.batch_size = 6 + infer.b_req_idx = torch.tensor([0, 0, 0, 1, 1, 1], dtype=torch.int32) + infer.att_state_manager = resources + infer.hidden_collector = NoopHiddenCollector() + backend = SimpleNamespace( + model=SimpleNamespace(is_mtp_draft_model=False, mtp_manager=SimpleNamespace(get_decode_draft_step=lambda _: 2)), + uses_dynamic_spec_verify_layout=lambda: False, + ) + att = GDNLinearAttDecodeAttState(backend=backend, _infer_state_weak_ref=weakref.ref(infer)) + infer.decode_att_state1 = att + att.init_state() + assert len(att.forwards) == 2 + assert att.get_forward(0) is att.get_forward(1) + assert att.get_forward(3).resource is second + configs = [SSMRunConfig(8, 1), SSMRunConfig(32, 2)] + att.get_forward(0).run_config, att.get_forward(3).run_config = configs + output = TpPartBaseModel._create_model_output(None, PostLayerOutput(torch.ones(6, 1)), infer) + assert {c.state_name: c.payload for c in output.state_commits} == { + first.state_name: configs[0], + second.state_name: configs[1], + } + + +@pytest.mark.parametrize("kda", [False, True]) +@pytest.mark.parametrize("speculative", [False, True]) +def test_native_forward_preserves_algorithm_inputs(monkeypatch, kda, speculative): + from importlib import import_module + + owner, _, _ = make_state(KDAState if kda else GDNState, step=2 if speculative else 0, device="cpu") + module, function = ( + ("fla.ops.kda_decode", "fused_recurrent_kda") + if kda + else ("mtp_fused_recurrent", "mtp_fused_recurrent_gated_delta_rule") + if speculative + else ("fla.ops", "fused_recurrent_gated_delta_rule") + ) + calls = [] + + def recurrent(**kwargs): + calls.append(kwargs) + return kwargs["v"], None + + monkeypatch.setattr( + import_module("lightllm.common.basemodel.triton_kernel.linear_att." + module), function, recurrent + ) + reqs = torch.tensor([1, 2], dtype=torch.int32) + cu = torch.tensor([0, 3, 6], dtype=torch.int32) if speculative else None + forward = owner.prepare_decode(reqs, cu) + shape = (1, 6) if speculative else (2, 1) + c = owner.config + q = torch.zeros(*shape, c.num_k_heads, c.key_dim) + v = torch.zeros(*shape, c.num_v_heads, c.value_dim) + gate = torch.zeros(q.numel() // c.key_dim, c.num_v_heads * (c.key_dim if kda else 1)) + beta = torch.zeros(gate.shape[0], c.num_v_heads) + a_log, bias = torch.zeros(c.num_v_heads), torch.zeros(gate.shape[-1]) + assert forward.run_layer(1, q, q, v, gate, beta, a_log, bias) is v + call = calls[0] + assert call["initial_state"].data_ptr() == owner.ssm_state[1].data_ptr() + assert call["ssm_state_indices"] is forward.state_indices + if speculative: + assert torch.equal(call["cu_seqlens"], cu) + assert call["num_accepted_tokens"] is forward.accepted_token_counts + if kda: + assert call["raw_gate"].shape == (*shape, c.num_v_heads * c.key_dim) + assert call["raw_beta"].shape == (*shape, c.num_v_heads) + assert call["lower_bound"] == -5.0 + else: + assert call["a_raw"] is gate and call["b_raw"] is beta + assert forward.finish() is None diff --git a/unit_tests/models/glm5_next/test_cache.py b/unit_tests/models/glm5_next/test_cache.py index f07e853f20..fa39331c26 100644 --- a/unit_tests/models/glm5_next/test_cache.py +++ b/unit_tests/models/glm5_next/test_cache.py @@ -149,13 +149,13 @@ def test_hybrid_checkpoint_restore_and_packed_kv_copy( conv = linear.conv_state[:, 0, ..., :3].clone() ssm = linear.ssm_state[:, 0].clone() if mode == "replay": - linear.ssm_backend.raw_keys.zero_() - linear.ssm_backend.raw_values.zero_() - linear.ssm_backend.gates.fill_(-0.5) - linear.ssm_backend.betas.zero_() - linear.ssm_backend.cursors[0] = 1 - ssm = linear.ssm_backend.get_accepted_state(0).clone() - linear.ssm_backend.cursors[2] = 3 + linear._ssm.raw_keys.zero_() + linear._ssm.raw_values.zero_() + linear._ssm.gates.fill_(-0.5) + linear._ssm.betas.zero_() + linear._ssm.cursors[0] = 1 + ssm = linear._ssm.get_accepted_state(0).clone() + linear._ssm.cursors[2] = 3 if mtp_step: linear.mtp_state_index[2] = mtp_step resources.save_bounded_att_runtime_to_pages( @@ -179,7 +179,7 @@ def test_hybrid_checkpoint_restore_and_packed_kv_copy( assert torch.equal(linear.conv_state[:, 2, ..., :3], conv) assert torch.equal(linear.ssm_state[:, 2 * linear.ssm_slots_per_req], ssm) if mode == "replay": - assert linear.ssm_backend.cursors[2] == 0 + assert linear._ssm.cursors[2] == 0 if mtp_step: assert linear.mtp_state_index[2] == 0 # 整组边界不依赖旧窗口，恢复大小页无需改动无重叠窗口。 diff --git a/unit_tests/models/glm5_next/test_mtp.py b/unit_tests/models/glm5_next/test_mtp.py index 2aa33161f5..01c899c497 100644 --- a/unit_tests/models/glm5_next/test_mtp.py +++ b/unit_tests/models/glm5_next/test_mtp.py @@ -256,12 +256,12 @@ def test_kda_backend_routes_accepted_conv_and_ssm_states(dynamic, role: str, mod accepted = torch.zeros(4, dtype=torch.int32, device="cuda") accepted[ids] = torch.tensor(lengths, dtype=torch.int32, device="cuda") - 1 if mode == "compact": - selected.ssm_backend.accept_updates(ids, accepted, run_config=state.run_config) + selected._ssm.accept_updates(ids, accepted, run_config=state.get_forward(0).run_config) else: - selected.ssm_backend.accept_updates(ids, accepted) + selected._ssm.accept_updates(ids, accepted) for req, length in zip(requests, lengths): torch.testing.assert_close( - selected.ssm_backend.get_accepted_state(req)[0], + selected._ssm.get_accepted_state(req)[0], expected_ssm[req * width + length - 1], atol=2e-3, rtol=1e-2, diff --git a/unit_tests/models/glm5_next/test_pd_cache.py b/unit_tests/models/glm5_next/test_pd_cache.py index 0a9c871a75..40f64b2639 100644 --- a/unit_tests/models/glm5_next/test_pd_cache.py +++ b/unit_tests/models/glm5_next/test_pd_cache.py @@ -169,7 +169,7 @@ def test_pd_kv_and_runtime_state_roundtrip( mem.linear.conv_state[:, src_req, ..., :3].copy_(conv[:, :, heads].reshape(6, -1, 3)) mem.linear.ssm_state[:, src_req * mem.linear.ssm_slots_per_req].copy_(ssm[:, heads]) if source_mode == "replay": - backend = mem.linear.ssm_backend + backend = mem.linear._ssm backend.raw_keys.zero_() backend.raw_values.zero_() backend.gates.fill_(-0.5) @@ -190,7 +190,7 @@ def test_pd_kv_and_runtime_state_roundtrip( mem.linear.conv_state.normal_() mem.linear.ssm_state.normal_() if dest_mode == "replay": - mem.linear.ssm_backend.cursors[dst_req] = 3 + mem.linear._ssm.cursors[dst_req] = 3 if mem.linear.mtp_state_index is not None: mem.linear.mtp_state_index[dst_req] = decode_mtp_step for state in mem.tails: @@ -262,7 +262,7 @@ def test_pd_kv_and_runtime_state_roundtrip( expected_tail[:, position % (4 + decode_mtp_step)] = tail[:, position % tail.shape[1]] assert torch.equal(restored_tail, expected_tail) if dest_mode == "replay": - assert mem.linear.ssm_backend.cursors[dst_req] == 0 + assert mem.linear._ssm.cursors[dst_req] == 0 if mem.linear.mtp_state_index is not None: assert mem.linear.mtp_state_index[dst_req] == 0 for (state, stride), expected in zip( From 26627bf1d537cdb482d0fb5bd62a48a3c2497cf6 Mon Sep 17 00:00:00 2001 From: sufubao Date: Tue, 3 Nov 2026 15:08:24 +0800 Subject: [PATCH 38/40] fix: preserve logical GDN layer state in loop models --- docs/replay_ssm.md | 4 + .../common/basemodel/attention/linear/gdn.py | 10 +- .../models/sensenova_loop/test_gdn_state.py | 150 ++++++++++++++++++ 3 files changed, 162 insertions(+), 2 deletions(-) create mode 100644 unit_tests/models/sensenova_loop/test_gdn_state.py diff --git a/docs/replay_ssm.md b/docs/replay_ssm.md index c074a15ea5..b78b06fbad 100644 --- a/docs/replay_ssm.md +++ b/docs/replay_ssm.md @@ -52,6 +52,10 @@ The execution path is: the existing fused accepted-position kernel; Compact reconstructs each partition using its own forward's launch config. +Weights use checkpoint (physical) layer IDs. Conv, SSM and KV state use logical +execution layer IDs throughout dispatch. Loop models may reuse one weight object +at multiple execution depths; its `layer_num_` must never select recurrent state. + Logical request arrays are `[B]`; verification rows are `[T]`. Partition sizes count real requests, excluding graph padding. Zero-sized partitions are allowed. Registered resource names include the model role, so a draft receipt cannot be diff --git a/lightllm/common/basemodel/attention/linear/gdn.py b/lightllm/common/basemodel/attention/linear/gdn.py index 2018f96031..0e44139037 100644 --- a/lightllm/common/basemodel/attention/linear/gdn.py +++ b/lightllm/common/basemodel/attention/linear/gdn.py @@ -229,6 +229,7 @@ def decode_att( b, self.infer_state, layer_weight, + layer_num=layer_num, ) else: # 非 MTP 模式下，使用线性层 decode 状态。 @@ -241,6 +242,7 @@ def decode_att( b, self.infer_state, layer_weight, + layer_num=layer_num, ) return core_attn_out, z @@ -254,6 +256,8 @@ def _gdn_decode_kernel( b: torch.Tensor, infer_state: "Qwen3NextInferStateInfo", layer_weight: "Qwen3NextTransformerLayerWeight", + *, + layer_num: int, ): backend: GDNLinearAttBackend = self.backend @@ -277,7 +281,7 @@ def _gdn_decode_kernel( backend.head_v_dim, ) core_attn_out = forward.run_layer( - layer_weight.layer_num_, + layer_num, query, key, value, @@ -297,6 +301,8 @@ def _gdn_mtp_kernel( b: torch.Tensor, infer_state: "Qwen3NextInferStateInfo", layer_weight: "Qwen3NextTransformerLayerWeight", + *, + layer_num: int, ): from lightllm.common.basemodel.triton_kernel.linear_att.causal_conv1d_mtp import ( causal_conv1d_update as causal_conv1d_update_mtp, @@ -319,7 +325,7 @@ def _gdn_mtp_kernel( query, key, value = backend._rearrange_mixed_qkv(mixed_qkv, decode=False) core_attn_out = forward.run_layer( - layer_weight.layer_num_, + layer_num, query, key, value, diff --git a/unit_tests/models/sensenova_loop/test_gdn_state.py b/unit_tests/models/sensenova_loop/test_gdn_state.py new file mode 100644 index 0000000000..a31c0acaf8 --- /dev/null +++ b/unit_tests/models/sensenova_loop/test_gdn_state.py @@ -0,0 +1,150 @@ +from types import SimpleNamespace + +import pytest +import torch + +from lightllm.common.attention_state import AttentionStateManager +from lightllm.common.attention_state.bounded_att import ConvSSMStateConfig, GDNState +from lightllm.common.basemodel.attention.base_att import AttControl +from lightllm.common.basemodel.attention.linear.triton import TritonLinearAttBackend +from lightllm.common.basemodel.infer_struct import InferStateInfo +from lightllm.models.sensenova_loop.config import get_loop_layer_indices, get_loop_layer_types +from lightllm.server.core.objs.start_args_type import StartArgs + + +@pytest.mark.parametrize("device", ["cpu", "cuda"]) +@pytest.mark.parametrize("mtp_step", [0, 2]) +@pytest.mark.parametrize( + "start,end,count,logical_layer,storage_layer", + [(4, 28, 1, 4, 3), (4, 28, 3, 28, 21), (1, 6, 2, 13, 10)], + ids=["no-loop", "aligned-loop", "unaligned-loop"], +) +def test_gdn_decode_keeps_logical_layer_state( + monkeypatch, device, mtp_step, start, end, count, logical_layer, storage_layer +): + if device == "cuda" and not torch.cuda.is_available(): + pytest.skip("CUDA required") + torch.manual_seed(1585) + settings = dict(start=start, end=end, eval_count=count, use_gate=False) + config = dict( + num_hidden_layers=32, + full_attention_interval=4, + layer_types=["full_attention" if (i + 1) % 4 == 0 else "linear_attention" for i in range(32)], + linear_num_key_heads=1, + linear_num_value_heads=2, + linear_key_head_dim=32, + linear_value_head_dim=32, + linear_conv_kernel_dim=4, + hidden_act="silu", + ) + indices = get_loop_layer_indices(settings, 32) + kinds = get_loop_layer_types(settings, config) + layers = tuple(i for i, kind in enumerate(kinds) if kind == "linear_attention") + args = StartArgs(mtp_step=mtp_step, linear_att_ssm_data_type="float32") + monkeypatch.setattr("lightllm.common.basemodel.attention.linear.gdn.get_env_start_args", lambda: args) + monkeypatch.setattr("lightllm.common.basemodel.attention.base_att.get_env_start_args", lambda: args) + backend = object.__new__(TritonLinearAttBackend) + backend.__init__( + SimpleNamespace( + config=config, + tp_world_size_=1, + is_mtp_draft_model=False, + mtp_manager=SimpleNamespace(get_decode_draft_step=lambda _: mtp_step), + ) + ) + dtype = torch.bfloat16 if device == "cuda" else torch.float32 + layout = ConvSSMStateConfig(layers, 1, 2, 32, 32, 3, dtype, torch.float32) + resource = GDNState(layout, request_capacity=3, mtp_step=mtp_step, device=device) + resources = AttentionStateManager(request_capacity=3) + resources.register_state("main", resource) + resources.finish_state_registration() + assert resource.get_storage_layer_index(logical_layer) == storage_layer + physical_weights = [ + SimpleNamespace( + layer_num_=i, + linear_conv1d=SimpleNamespace( + mm_param=SimpleNamespace(weight=torch.randn(layout.conv_channels, 4, device=device, dtype=dtype) * 0.1), + bias=None, + ), + linear_A_log=SimpleNamespace(weight=torch.zeros(2, device=device)), + linear_dt_bias=SimpleNamespace(weight=torch.zeros(2, device=device)), + ) + for i in range(32) + ] + execution_weights = [physical_weights[i] for i in indices] + weight = execution_weights[logical_layer] + assert weight is physical_weights[indices[logical_layer]] + + if device == "cpu": + + def conv_pack(mixed, z, a, b, conv, *args): + conv[1].add_(1) + return (*backend._rearrange_mixed_qkv(mixed, decode=True), z, a, b) + + def conv_mtp(mixed, conv, *args, **kwargs): + conv[1].add_(1) + return mixed + + def recurrent(*, initial_state, v, **kwargs): + initial_state[mtp_step + 1 : 2 * (mtp_step + 1)].add_(1) + return v, None + + monkeypatch.setattr("lightllm.common.basemodel.attention.linear.gdn.conv_pack_gdn_decode_inputs", conv_pack) + monkeypatch.setattr( + "lightllm.common.basemodel.triton_kernel.linear_att.causal_conv1d_mtp.causal_conv1d_update", conv_mtp + ) + monkeypatch.setattr( + "lightllm.common.basemodel.triton_kernel.linear_att.fla.ops.fused_recurrent_gated_delta_rule", recurrent + ) + monkeypatch.setattr( + "lightllm.common.basemodel.triton_kernel.linear_att.mtp_fused_recurrent." + "mtp_fused_recurrent_gated_delta_rule", + recurrent, + ) + + infer = InferStateInfo() + infer.att_state_manager, infer.att_state_role = resources, "main" + infer.b_req_idx = torch.tensor([1], device=device, dtype=torch.int32) + control = dict(layer_num=logical_layer, layer_weight=weight) + mixed_dim = backend.tp_qkvz_dim + backend.tp_ba_dim + if device == "cuda": + infer.b1_cu_q_seq_len = torch.tensor([0, 5], device=device, dtype=torch.int32) + infer.b_ready_cache_len = torch.zeros(1, device=device, dtype=torch.int32) + prefill = backend.create_att_prefill_state(infer) + prefill.init_state() + prefill.prefill_att( + None, + None, + None, + AttControl( + linear_att_prefill=True, + linear_att_prefill_dict=dict( + **control, mixed_qkvzba=torch.randn(5, mixed_dim, device=device, dtype=dtype) + ), + ), + ) + + infer.batch_size = mtp_step + 1 + infer.b_req_idx = infer.b_req_idx.repeat(infer.batch_size) + decode = backend.create_att_decode_state(infer) + untouched = [i for i in range(len(layers)) if i != storage_layer] + for _ in range(2): + before_conv, before_ssm = resource.conv_state.clone(), resource.ssm_state.clone() + decode.init_state() + output, _ = decode.decode_att( + None, + None, + None, + AttControl( + linear_att_decode=True, + linear_att_decode_dict=dict( + **control, mixed_qkvzba=torch.randn(infer.batch_size, mixed_dim, device=device, dtype=dtype) + ), + ), + ) + assert torch.isfinite(output).all() + assert not torch.equal(resource.conv_state[storage_layer], before_conv[storage_layer]) + assert not torch.equal(resource.ssm_state[storage_layer], before_ssm[storage_layer]) + assert torch.equal(resource.conv_state[untouched], before_conv[untouched]) + assert torch.equal(resource.ssm_state[untouched], before_ssm[untouched]) + assert weight.layer_num_ == indices[logical_layer] From 38d9a6bed5dc872487887acb427bcf0e2cd0d5ea Mon Sep 17 00:00:00 2001 From: sufubao Date: Tue, 3 Nov 2026 15:08:34 +0800 Subject: [PATCH 39/40] test: initialize runtime args for SSM checkpoint tests --- .../triton_kernel/linear_att/test_replayssm.py | 2 +- unit_tests/common/test_linear_state.py | 4 +++- unit_tests/conftest.py | 13 +++++++++++++ 3 files changed, 17 insertions(+), 2 deletions(-) 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 index 4713d1015a..db1ab6d98e 100644 --- a/unit_tests/common/basemodel/triton_kernel/linear_att/test_replayssm.py +++ b/unit_tests/common/basemodel/triton_kernel/linear_att/test_replayssm.py @@ -549,7 +549,7 @@ def test_autotune_candidates_match_recurrence(dtype, run_config, mode): @pytest.mark.parametrize("save_big", [False, True]) @pytest.mark.parametrize("dtype", [torch.float32, torch.bfloat16]) -def test_checkpoint_restores_pending_history_and_accepted_conv(save_big, dtype): +def test_checkpoint_restores_pending_history_and_accepted_conv(runtime_start_args, save_big, dtype): from lightllm.common.attention_state.bounded_att import ConvSSMStateConfig from unit_tests.common.test_linear_state import make_state diff --git a/unit_tests/common/test_linear_state.py b/unit_tests/common/test_linear_state.py index 25f00f42e3..8eb25c073d 100644 --- a/unit_tests/common/test_linear_state.py +++ b/unit_tests/common/test_linear_state.py @@ -31,7 +31,9 @@ def make_state(cls=GDNState, mode="native", step=2, dtype=torch.float32, device= @pytest.mark.parametrize("cls", [GDNState, KDAState]) @pytest.mark.parametrize("mode,step", [("native", 0), ("native", 2), ("compact", 2), ("replay", 0), ("replay", 2)]) @pytest.mark.parametrize("dtype", [torch.float32, torch.bfloat16]) -def test_checkpoint_and_pd_export_same_accepted_prefix_and_restore_reused_slot(cls, mode, step, dtype): +def test_checkpoint_and_pd_export_same_accepted_prefix_and_restore_reused_slot( + runtime_start_args, cls, mode, step, dtype +): state, _, layout = make_state(cls, mode, step, dtype) state.initialize_page_pools(big_page_capacity=2, small_page_capacity=2) src, dst, offset = 1, 2, min(1, step) diff --git a/unit_tests/conftest.py b/unit_tests/conftest.py index c07fd123f0..20d59199bc 100644 --- a/unit_tests/conftest.py +++ b/unit_tests/conftest.py @@ -9,6 +9,19 @@ import pytest +@pytest.fixture +def runtime_start_args(monkeypatch: pytest.MonkeyPatch): + import json + from dataclasses import asdict + from lightllm.server.core.objs.start_args_type import StartArgs + from lightllm.utils.envs_utils import get_env_start_args + + monkeypatch.setenv("LIGHTLLM_START_ARGS", json.dumps(asdict(StartArgs()))) + get_env_start_args.cache_clear() + yield + get_env_start_args.cache_clear() + + @pytest.fixture def kv_allocator_environment(monkeypatch: pytest.MonkeyPatch): """Provide a service identity and reclaim the allocator's real shared counters.""" From cfc9f8bfecc45a650065706374b86172ec4fa0f3 Mon Sep 17 00:00:00 2001 From: sufubao Date: Tue, 3 Nov 2026 15:15:42 +0800 Subject: [PATCH 40/40] docs: remove ReplaySSM design document --- docs/replay_ssm.md | 128 --------------------------------------------- 1 file changed, 128 deletions(-) delete mode 100644 docs/replay_ssm.md diff --git a/docs/replay_ssm.md b/docs/replay_ssm.md deleted file mode 100644 index b78b06fbad..0000000000 --- a/docs/replay_ssm.md +++ /dev/null @@ -1,128 +0,0 @@ -# Recurrent state ownership - -GDN and KDA share the Conv/SSM resource lifecycle. Their attention implementations -own convolution, input packing and prefill computation. Each resource composes -one SSM implementation that owns recurrent execution, speculative storage and -accepted-state reconstruction. - -```mermaid -flowchart TD - Runner[Model runner: request layout and stream ordering] --> Attention[GDN / KDA attention] - Attention --> Forward[SSMForward: one resource and one invocation] - Forward --> Resource[ConvSSMState: accepted prefix and canonical snapshots] - Resource --> Native[NativeSSMState: candidate slots] - Resource --> Compact[CompactSSMState: verify scratch and reconstruction] - Resource --> Replay[ReplaySSMState: history and folding] - Native --> Kernels[Triton kernels] - Compact --> Kernels - Replay --> Kernels - Forward --> Output[ModelOutput.state_commits: immutable resource-keyed receipts] - Output --> Verify[MTP verification and logical microbatch partitions] - Verify --> Manager[AttentionStateManager: validate and dispatch] - Manager --> Resource - Resource --> Transfer[Canonical Conv/SSM pages: CPU cache / DP / PD] -``` - -## Lifetimes and interfaces - -| Owner | Lifetime | Responsibility | -|---|---|---| -| `ConvSSMState` | Worker/resource | Conv pool, accepted position, layer mapping, canonical save/restore | -| Private `SSMState` implementation | Worker/resource | SSM pool, candidate/history layout, kernels, acceptance | -| `SSMForward` | Forward or captured graph | Request/state indices, cumulative lengths, prior accepted counts, selected launch config | -| `StateCommit` | Model output or captured graph | Registered resource name and immutable implementation metadata | -| `StateCommitBatch` | Verification call | Number of real logical requests and receipts from their producing forward | - -The execution path is: - -1. `LinearDecodeAttState.init_state()` groups fixed or dynamic MTP rows once and - asks each participating resource to `prepare_decode()`. GDN obtains its actual - draft width from `MtpManager`; KDA retains its configured width. -2. GDN/KDA select the resource by role, algorithm and logical layer. After their - existing convolution/packing, `forward.run_layer()` dispatches the recurrence. - The attention caller does not inspect storage modes, cursors or kernel configs. -3. `forward.finish()` publishes a detached receipt when acceptance needs metadata. - Compact publishes the exact `SSMRunConfig` used for verification. Native and - Replay need no receipt. Multiple resources can publish independently. -4. `verify_mtp_tokens()` passes logical request IDs and accepted lengths to the - generic bounded-state commit interface. `MtpCommitContext` also retains the - expanded verification rows and partitions from the producing forwards. -5. The manager validates all receipts before dispatching any mutations. Each SSM - implementation commits its own state. Replay supplies its private cursor to - the existing fused accepted-position kernel; Compact reconstructs each - partition using its own forward's launch config. - -Weights use checkpoint (physical) layer IDs. Conv, SSM and KV state use logical -execution layer IDs throughout dispatch. Loop models may reuse one weight object -at multiple execution depths; its `layer_num_` must never select recurrent state. - -Logical request arrays are `[B]`; verification rows are `[T]`. Partition sizes -count real requests, excluding graph padding. Zero-sized partitions are allowed. -Registered resource names include the model role, so a draft receipt cannot be -applied to the main model. This uses the existing registration identity; it does -not add arbitrary same-type resource groups to the registration/transport schema. - -Missing Compact receipts, unknown resource names, duplicate receipts and invalid -partition totals are errors. No fallback reads a mutable "last forward" config. -The runner must commit exactly once, preserve request ordering, and prevent -overlapping live forwards from using the same real request slots. A receipt is -reusable across graph invocations; it is not a one-use transaction or rollback -mechanism. Tensor-value checks requiring a device synchronization are not added -to the decode path. - -## CUDA Graph and microbatch rules - -Request preparation runs before every invocation, including graph replay. -Replay snapshots its forward cursor before accepting the guaranteed token of -ordinary decode. Moving preparation solely into capture-time Python would break -later invocations. - -`SSMForward.copy_inputs()` copies nested tensor metadata into captured buffers, -validates their layouts, and preserves the captured launch config. Aliased -request/state-index buffers are copied once. Warmup clones each forward wrapper -so configuration selection cannot mutate the wrapper used for capture. Storage -and the tuning cache remain shared intentionally. - -Graph replay returns the captured `ModelOutput`; unpadding preserves its receipt -tuple. Overlapped microbatches pass separate `StateCommitBatch` entries, each with -its own real request count and configuration. Neither the scheduler nor the -generic manager needs to know that Compact uses those configs for reconstruction. - -## Checkpoints and numerical behavior - -`get_accepted_state()` returns the convolution window and SSM state at the same -accepted prefix. Native selects a candidate slot, Compact exposes its committed -state, and Replay materializes a snapshot without folding or changing live -history. CPU pages and DP/PD transport contain this canonical state, independent -of execution mode or speculative width. `prepare_restore()` clears history and -the accepted position before exposing destinations for both fields. - -Views may alias live storage. Existing stream ordering and completion events must -protect consuming copies before advancing or recycling requests. Restore callers -must fill both fields before execution resumes. A token length labels the -accepted prefix; it does not request reconstruction of arbitrary earlier tokens. - -Native and Compact retain per-token state rounding. Replay retains its FP32 -history accumulation and rounding at fold/materialization. These are intentionally -different numerical semantics. The refactor relocates host orchestration while -preserving the Triton recurrence/folding kernels and their launch parameters. -CLI modes, replay configuration, grouped heads, KDA gate bounds and existing -parallel-mode restrictions are unchanged. - -## Source map and verification - -- `common/attention_state/bounded_att/conv_ssm_state.py`: shared resource and - canonical page/transfer lifecycle; GDN/KDA subclasses supply identity/page type. -- `common/attention_state/bounded_att/ssm/`: native, compact, replay and forward - lifetimes; only the factory dispatches the configured mode. -- `common/basemodel/attention/linear/state.py`: shared request layout and graph - integration; `gdn.py`/`kda.py` retain algorithm-specific processing. -- `common/attention_state/commit.py`: generic verification and receipt contracts. -- `common/basemodel/triton_kernel/linear_att/`: numerical kernels and tuning. - -CPU contract tests cover accepted-state selection/reuse, resource-key routing, -invalid receipts, per-microbatch configs, nested graph copying and warmup -isolation. Kernel substitutions in these tests verify orchestration only. -The retained CUDA tests cover numerical recurrence, fused acceptance, dynamic -layouts, folding, graph alternation, checkpoint continuation and IPC. They must -run on compatible GPU hardware before claiming GPU or serving acceptance.