iHore/vLLM-4x-RTX-3090
8
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity
Files
main
vLLM-4x-RTX-3090/patches/0080-recoverssm-allow-pipeline-parallelism-for-Qwen4Exp-o.patch
T
MOVZX 777e910dd0 Initial commit 
Signed-off-by: MOVZX <movzx@yahoo.com>
2026-10-15 01:01:52 +07:00

1018 lines
42 KiB
Diff
Raw Permalink Blame History

From 9917a9883e710ddae023d31392c0d5804d2b3b6c Mon Sep 17 00:00:00 2001
From: halt95 <halt95@users.noreply.github.com>
Date: Thu, 1 Oct 2026 21:15:44 -0400
Subject: [PATCH 080/103] recoverssm: allow pipeline parallelism for Qwen4Exp
on the V2 runner
Upstream RecoverSSM commits a step in the same step's postprocess, hence PP=1. Under pipeline parallelism a
non-last rank learns a step's accepted counts only when the PP broadcast arrives pp_size steps later, while later
builds overwrite the persistent commit inputs. The commit is now deferred with its own copy of those inputs, and
the PP=1 check is lifted for Qwen4Exp on Model Runner V2. Config and CPU tests cover the deferral.
---
tests/config/test_recoverssm_config.py | 13 +-
tests/config/test_recoverssm_pp_config.py | 55 +++
.../worker/gpu/test_recoverssm_pp_deferral.py | 408 ++++++++++++++++++
vllm/config/vllm.py | 22 +-
vllm/v1/attention/backends/gdn_recoverssm.py | 34 +-
vllm/v1/attention/backends/ple_recoverssm.py | 23 +-
.../attention/backends/recoverssm_metadata.py | 41 +-
vllm/v1/worker/gpu/cudagraph_utils.py | 4 +
vllm/v1/worker/gpu/model_runner.py | 11 +
vllm/v1/worker/gpu/model_states/interface.py | 13 +
.../worker/gpu/model_states/mamba_hybrid.py | 14 +
vllm/v1/worker/gpu/model_states/recoverssm.py | 68 +++
vllm/v1/worker/gpu/pp_utils.py | 31 +-
13 files changed, 727 insertions(+), 10 deletions(-)
create mode 100644 tests/config/test_recoverssm_pp_config.py
create mode 100644 tests/v1/worker/gpu/test_recoverssm_pp_deferral.py
diff --git a/tests/config/test_recoverssm_config.py b/tests/config/test_recoverssm_config.py
index 27e26d9..2fba7eb 100644
--- a/tests/config/test_recoverssm_config.py
+++ b/tests/config/test_recoverssm_config.py
@@ -83,19 +83,26 @@ def test_derived_drafter_keeps_the_choice_and_runs_the_checks():
drafter = _config(DRAFTER, cache_config=target.cache_config)
_validator()(drafter)
assert target.cache_config.use_recoverssm is True
- bad = _config(
+ # Fork: PP > 1 is lifted for Qwen4Exp on Model Runner V2; the
+ # drafter's derived config runs the same check and passes it.
+ pp2 = _config(
DRAFTER,
cache_config=target.cache_config,
parallel_config__pipeline_parallel_size=2,
)
- with pytest.raises(ValueError, match="pipeline_parallel_size"):
+ _validator()(pp2)
+ bad = _config(
+ DRAFTER,
+ cache_config=target.cache_config,
+ mamba_config__backend="flashinfer",
+ )
+ with pytest.raises(ValueError, match="mamba-backend triton"):
_validator()(bad)
@pytest.mark.parametrize(
("override", "value", "match"),
[
- ("parallel_config__pipeline_parallel_size", 2, "pipeline_parallel_size"),
("mamba_config__enable_stochastic_rounding", True, "stochastic"),
("mamba_config__backend", "flashinfer", "mamba-backend triton"),
("use_v2_model_runner", False, "V2_MODEL_RUNNER"),
diff --git a/tests/config/test_recoverssm_pp_config.py b/tests/config/test_recoverssm_pp_config.py
new file mode 100644
index 0000000..95117f9
--- /dev/null
+++ b/tests/config/test_recoverssm_pp_config.py
@@ -0,0 +1,55 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""The RecoverSSM pipeline-parallel lift at the config level.
+
+Upstream (#58863) rejects RecoverSSM for pipeline_parallel_size > 1. The fork
+lifts that for Qwen4Exp GDN/PLE on Model Runner V2, where non-last PP ranks
+defer the commit with a snapshot of its inputs. Kimi-K3 KDA keeps the PP=1
+check (its metadata has no snapshot), and Model Runner V1 has no deferral.
+
+Failing control: on the unlifted port (upstream check), test_pp2_qwen4exp_v2_*
+raises "RecoverSSM currently requires pipeline_parallel_size=1".
+"""
+
+import pytest
+
+from tests.config.test_recoverssm_config import _config, _validator
+
+TARGET = "Qwen4ExpForConditionalGeneration"
+
+
+@pytest.mark.parametrize("pp", [2, 4])
+@pytest.mark.parametrize(
+ "arch", ["Qwen4ExpForConditionalGeneration", "Qwen4ExpForCausalLM", "Qwen4ExpMTP"]
+)
+def test_pp2_qwen4exp_v2_is_accepted(arch, pp):
+ cfg = _config(arch, parallel_config__pipeline_parallel_size=pp)
+ _validator()(cfg)
+ assert cfg.cache_config.use_recoverssm is True
+
+
+def test_pp2_kda_keeps_the_pp1_requirement():
+ cfg = _config(
+ "KimiK3ForConditionalGeneration", parallel_config__pipeline_parallel_size=2
+ )
+ with pytest.raises(ValueError, match="only for Qwen4Exp"):
+ _validator()(cfg)
+
+
+def test_pp2_v1_runner_is_rejected():
+ # none mode, so the align-mode V2 check does not fire first
+ cfg = _config(TARGET, parallel_config__pipeline_parallel_size=2)
+ cfg.cache_config.mamba_cache_mode = "none"
+ cfg.use_v2_model_runner = False
+ with pytest.raises(ValueError, match="pipeline_parallel_size > 1 requires"):
+ _validator()(cfg)
+
+
+def test_pp2_still_runs_the_other_checks():
+ cfg = _config(
+ TARGET,
+ parallel_config__pipeline_parallel_size=2,
+ mamba_config__backend="flashinfer",
+ )
+ with pytest.raises(ValueError, match="mamba-backend triton"):
+ _validator()(cfg)
diff --git a/tests/v1/worker/gpu/test_recoverssm_pp_deferral.py b/tests/v1/worker/gpu/test_recoverssm_pp_deferral.py
new file mode 100644
index 0000000..9420aaa
--- /dev/null
+++ b/tests/v1/worker/gpu/test_recoverssm_pp_deferral.py
@@ -0,0 +1,408 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""The RecoverSSM commit under pipeline parallelism (CPU, no GPU).
+
+A non-last PP rank verifies a speculative window at step T but receives the
+accepted counts only when the PP broadcast of step T is consumed, pp_size steps
+later (``PPHandler.get_prev_sampled_outputs`` in ``update_pp_decode_requests``).
+Upstream #58863 keeps one recorded step (``RecoverSSMState._step``) and commits
+it in the same step's postprocess, so it requires pipeline_parallel_size=1:
+under PP the deferred postprocess of step T would commit the metadata of the
+most recent step instead, and that metadata's spec state indices, query offsets
+and batch-order block table are persistent buffers the next build overwrites.
+
+The fork detaches the step in ``sample_tokens`` (a snapshot whose commit inputs
+are clones), keeps it on step T's pending PP entry, restores it right before
+step T's deferred postprocess, and masks rows whose request was freed in
+between (their blocks may already belong to another request).
+
+Every scenario here has a failing control that flips: the upstream single-slot
+flow and a snapshot that aliases the persistent buffers both commit the wrong
+step; a deferred commit without the freed-row mask commits into a freed slot;
+a teardown without the PP purge keeps the profiling KV cache alive.
+"""
+
+import ast
+import gc
+import inspect
+import textwrap
+import weakref
+from collections import deque
+from dataclasses import dataclass, fields
+from types import SimpleNamespace
+
+import numpy as np
+import pytest
+import torch
+
+from vllm.v1.attention.backends import gdn_recoverssm as gdn_mod
+from vllm.v1.attention.backends.gdn_recoverssm import (
+ GDNRecoverSSMCommitMetadata,
+ GDNRecoverSSMMetadata,
+)
+from vllm.v1.attention.backends.ple_recoverssm import (
+ PleRecoverSSMCommitMetadata,
+ PleRecoverSSMMetadata,
+)
+from vllm.v1.worker.gpu import pp_utils as pp_utils_mod
+from vllm.v1.worker.gpu.model_states.recoverssm import RecoverSSMState
+
+MAX_REQS = 8
+SPEC_LEN = 4 # 1 + K, K = 3 (prod)
+
+
+class _RecordingContext:
+ """Stands in for GDNRecoverSSMCommitContext: records what each commit
+ reads, at commit time (clones, like the kernels reading device memory)."""
+
+ def __init__(self) -> None:
+ self.commits: list[dict] = []
+
+ def commit(self, num_accepted_tokens, state_indices, query_start_loc, **kw):
+ req = kw.get("request_indices")
+ rows = (
+ list(range(state_indices.shape[0]))
+ if req is None
+ else [int(r) for r in req.tolist()]
+ )
+ self.commits.append(
+ dict(
+ state_indices=state_indices.tolist(),
+ query_start_loc=query_start_loc.tolist(),
+ num_accepted=[
+ int(num_accepted_tokens[r])
+ if r < num_accepted_tokens.shape[0]
+ else None
+ for r in rows
+ ],
+ block_table=None
+ if kw.get("block_table") is None
+ else kw["block_table"].clone(),
+ )
+ )
+
+
+class _Buffers:
+ """The builder's persistent, graph-staged buffers (one set per builder)."""
+
+ def __init__(self) -> None:
+ self.spec_state_indices = torch.full((MAX_REQS, SPEC_LEN), -1, dtype=torch.int32)
+ self.spec_query_start_loc = torch.zeros(MAX_REQS + 1, dtype=torch.int32)
+ self.block_table = torch.zeros(MAX_REQS, 6, dtype=torch.int32)
+
+
+def _gdn_meta(buf: _Buffers, slots, ctx, *, align=False) -> GDNRecoverSSMMetadata:
+ """Build step metadata the way GDNRecoverSSMMetadataBuilder.build does: the
+ commit points at views of the persistent buffers, written in place."""
+ n = len(slots)
+ buf.spec_state_indices[:n, 0] = torch.tensor(slots, dtype=torch.int32)
+ buf.spec_query_start_loc[: n + 1] = torch.arange(
+ 0, (n + 1) * SPEC_LEN, SPEC_LEN, dtype=torch.int32
+ )
+ for i, s in enumerate(slots):
+ buf.block_table[i] = s * 10 + torch.arange(6, dtype=torch.int32)
+ commit = GDNRecoverSSMCommitMetadata(
+ state_indices=buf.spec_state_indices[:, :1],
+ query_start_loc=buf.spec_query_start_loc,
+ request_indices=None,
+ block_table=buf.block_table[:n] if align else None,
+ num_computed_tokens=torch.full((n,), 100, dtype=torch.int32)
+ if align
+ else None,
+ block_size=64 if align else None,
+ )
+ required = dict(
+ num_prefills=0,
+ num_prefill_tokens=0,
+ num_decodes=0,
+ num_decode_tokens=0,
+ num_spec_decodes=n,
+ num_spec_decode_tokens=n * SPEC_LEN,
+ num_actual_tokens=n * SPEC_LEN,
+ )
+ kwargs = {
+ f.name: required.get(f.name)
+ for f in fields(GDNRecoverSSMMetadata)
+ if f.init and f.name not in ("recoverssm_commit", "recoverssm_context")
+ }
+ return GDNRecoverSSMMetadata(
+ **kwargs, recoverssm_commit=commit, recoverssm_context=ctx
+ )
+
+
+def _record(state: RecoverSSMState, meta) -> None:
+ group = SimpleNamespace(layer_names=["model.layers.0.linear_attn"])
+ state.record_step(
+ {"model.layers.0.linear_attn": meta}, [[group]], for_capture=False
+ )
+
+
+def _handler() -> pp_utils_mod.PPHandler:
+ h = object.__new__(pp_utils_mod.PPHandler)
+ h.is_last_rank = False
+ h.device = torch.device("cpu")
+ h.main_stream = SimpleNamespace(wait_event=lambda event: None)
+ h.req_idx_gen_np = np.zeros(MAX_REQS, dtype=np.int32)
+ h.queue = deque([None, None]) # pp_size = 2, pre-seeded like __init__
+ return h
+
+
+def _receive(h: pp_utils_mod.PPHandler, idx_mapping, num_sampled) -> None:
+ """What PPHandler.receive leaves in the queue (minus the NCCL broadcast)."""
+ idx_np = np.asarray(idx_mapping, dtype=np.int64)
+ n = len(idx_mapping)
+ h.queue[-1] = pp_utils_mod.PendingRecv(
+ event=object(),
+ sampled_tokens=torch.zeros(n, SPEC_LEN, dtype=torch.int64),
+ num_sampled=torch.tensor(num_sampled, dtype=torch.int32),
+ num_rejected=torch.zeros(n, dtype=torch.int32),
+ idx_mapping=torch.tensor(idx_mapping, dtype=torch.int32),
+ idx_mapping_np=idx_np,
+ need_sampled_mask=np.ones(n, dtype=bool),
+ gen_at_receive_np=h.req_idx_gen_np[idx_np].copy(),
+ )
+
+
+@pytest.fixture
+def h2d(monkeypatch):
+ monkeypatch.setattr(
+ pp_utils_mod,
+ "async_tensor_h2d",
+ lambda a, device=None, dtype=None, out=None: torch.as_tensor(
+ np.asarray(a), dtype=torch.int32
+ ),
+ )
+
+
+def _run_two_microbatches(mode: str, *, free_slot: int | None = None):
+ """Non-last rank, pp_size 2: micro-batch A at step T (req slots 0, 1 with
+ state blocks 3, 5), micro-batch B at T+1 (req slots 2, 4, 6 with blocks 7,
+ 9, 11), then step T+2 consumes A's broadcast. Returns the commit A got.
+
+ mode: "fork" (detach/attach/restore), "upstream" (#58863 single slot) or
+ "alias" (fork flow, but the snapshot aliases the persistent buffers)."""
+ ctx = _RecordingContext()
+ buf = _Buffers() # one builder: both micro-batches share its buffers
+ state = RecoverSSMState()
+ h = _handler()
+
+ def sample_tokens(idx_mapping, num_sampled):
+ _receive(h, idx_mapping, num_sampled)
+ if mode != "upstream":
+ h.attach_deferred_state(state.detach_step())
+
+ # step T: micro-batch A (execute_model first consumes a pre-seeded None)
+ assert h.get_prev_sampled_outputs() is None
+ _record(state, _gdn_meta(buf, [3, 5], ctx))
+ sample_tokens([0, 1], [2, 4])
+ # step T+1: micro-batch B; its build overwrites the builder's buffers
+ assert h.get_prev_sampled_outputs() is None
+ _record(state, _gdn_meta(buf, [7, 9, 11], ctx))
+ sample_tokens([2, 4, 6], [1, 1, 1])
+ if free_slot is not None:
+ h.on_req_idx_freed(free_slot) # A's request freed before T+2
+ # step T+2: update_pp_decode_requests consumes A's entry
+ out = h.get_prev_sampled_outputs()
+ assert out is not None
+ deferred = out.pop("deferred_state")
+ if mode != "upstream":
+ state.restore_step(deferred)
+ state.commit_step(
+ out["num_sampled"],
+ out["idx_mapping"],
+ state_indices=None,
+ num_accepted_tokens=torch.ones(MAX_REQS, dtype=torch.int32),
+ )
+ assert ctx.commits, "no commit ran"
+ return ctx.commits[-1]
+
+
+A_COMMIT = dict(state_indices=[3, 5], num_accepted=[2, 4])
+
+
+def test_deferred_commit_commits_its_own_step(h2d):
+ got = _run_two_microbatches("fork")
+ assert got["state_indices"] == A_COMMIT["state_indices"]
+ assert got["num_accepted"] == A_COMMIT["num_accepted"]
+ assert got["query_start_loc"] == [0, 4, 8]
+
+
+def test_control_upstream_single_slot_commits_the_wrong_step(h2d):
+ # #58863 as-is under PP: step T's counts land on micro-batch B's blocks.
+ got = _run_two_microbatches("upstream")
+ assert got["state_indices"] != A_COMMIT["state_indices"]
+ assert got["state_indices"] == [7, 9, 11]
+
+
+def test_control_aliasing_snapshot_commits_the_wrong_blocks(h2d, monkeypatch):
+ # Detaching without cloning keeps views of the builder's persistent
+ # buffers: the next micro-batch's build rewrites them before the commit.
+ monkeypatch.setattr(
+ GDNRecoverSSMMetadata, "snapshot_for_deferred_commit", lambda self: self
+ )
+ got = _run_two_microbatches("alias")
+ assert got["state_indices"] != A_COMMIT["state_indices"]
+
+
+def test_freed_row_commits_nothing(h2d):
+ # Request slot 1 (block 5) freed between verify and commit: count 0 makes
+ # the commit plan skip it (commit_len 0), slot 0 still commits.
+ got = _run_two_microbatches("fork", free_slot=1)
+ assert got["state_indices"] == A_COMMIT["state_indices"]
+ assert got["num_accepted"] == [2, 0]
+
+
+def test_control_unmasked_deferred_commit_writes_a_freed_slot(h2d, monkeypatch):
+ real = RecoverSSMState.restore_step
+
+ def restore_without_mask(self, deferred):
+ real(self, deferred)
+ self._step_is_deferred = False # what a commit without the mask sees
+
+ monkeypatch.setattr(RecoverSSMState, "restore_step", restore_without_mask)
+ got = _run_two_microbatches("fork", free_slot=1)
+ assert got["num_accepted"] == [2, 4] # block 5 would be written
+
+
+def test_snapshot_clones_every_commit_input_gdn_align():
+ ctx = _RecordingContext()
+ buf = _Buffers()
+ meta = _gdn_meta(buf, [3, 5], ctx, align=True)
+ snap = meta.snapshot_for_deferred_commit()
+ before = (
+ snap.recoverssm_commit.state_indices.clone(),
+ snap.recoverssm_commit.query_start_loc.clone(),
+ snap.recoverssm_commit.block_table.clone(),
+ )
+ _gdn_meta(buf, [7, 9], ctx, align=True) # next build rewrites the buffers
+ c = snap.recoverssm_commit
+ assert torch.equal(c.state_indices, before[0])
+ assert torch.equal(c.query_start_loc, before[1])
+ assert torch.equal(c.block_table, before[2])
+ # control: the live metadata's views did change
+ assert not torch.equal(meta.recoverssm_commit.state_indices[:2], before[0])
+ assert not torch.equal(meta.recoverssm_commit.block_table, before[2])
+ for t in (c.state_indices, c.query_start_loc, c.block_table):
+ assert all(t.data_ptr() != b.data_ptr() for b in vars(buf).values())
+
+
+def test_snapshot_ple():
+ sidx = torch.tensor([3, 5, -1, -1], dtype=torch.int32)
+ qsl = torch.tensor([0, 4, 8, 8, 8], dtype=torch.int32)
+ commit = PleRecoverSSMCommitMetadata(
+ state_indices=sidx,
+ query_start_loc=qsl,
+ request_indices=None,
+ block_table=None,
+ num_computed_tokens=None,
+ block_size=None,
+ )
+ names = [
+ f.name
+ for f in fields(PleRecoverSSMMetadata)
+ if f.init and f.name not in ("recoverssm_commit", "recoverssm_context")
+ ]
+ meta = PleRecoverSSMMetadata(
+ **{n: None for n in names},
+ recoverssm_commit=commit,
+ recoverssm_context=_RecordingContext(),
+ )
+ meta.num_spec_decodes = 2
+ snap = meta.snapshot_for_deferred_commit()
+ sidx.fill_(9)
+ qsl.fill_(0)
+ assert snap.recoverssm_commit.state_indices.tolist() == [3, 5]
+ assert snap.recoverssm_commit.query_start_loc.tolist() == [0, 4, 8]
+
+
+def test_attach_without_a_pending_entry_raises():
+ h = _handler()
+ h.queue[-1] = None # receive() skipped the broadcast for this step
+ with pytest.raises(RuntimeError, match="without sampled outputs"):
+ h.attach_deferred_state(object())
+ h.attach_deferred_state(None) # nothing to keep: fine
+
+
+def test_restored_step_without_counts_raises():
+ state = RecoverSSMState()
+ ctx = _RecordingContext()
+ _record(state, _gdn_meta(_Buffers(), [3], ctx))
+ state.restore_step(state.detach_step())
+ with pytest.raises(RuntimeError, match="without sampled counts"):
+ state.commit_step(
+ 0, torch.zeros(1, dtype=torch.int32), state_indices=None,
+ num_accepted_tokens=torch.ones(1, dtype=torch.int32),
+ )
+
+
+def test_detach_of_a_step_without_spec_rows_is_none():
+ state = RecoverSSMState()
+ meta = _gdn_meta(_Buffers(), [3], _RecordingContext())
+ meta.recoverssm_commit = None
+ _record(state, meta)
+ assert state.detach_step() is None
+
+
+@dataclass
+class _Holder:
+ kv: torch.Tensor
+
+
+def test_teardown_drops_the_deferred_states_and_their_kv_cache():
+ """Profiling-state leak guard (the ~26 GiB on H200 class, #58863 commit 4):
+ a deferred step on a pending PP entry references the KV cache through its
+ commit context; the teardown must release it."""
+ h = _handler()
+ kv = torch.zeros(1024)
+ ref = weakref.ref(kv)
+ _receive(h, [0], [1])
+ h.attach_deferred_state(_Holder(kv))
+ del kv
+ gc.collect()
+ assert ref() is not None # control: held by the pending entry
+ assert h.drop_deferred_states() == 1
+ gc.collect()
+ assert ref() is None
+ assert h.queue[-1] is not None # the sampled-token entry itself stays
+
+
+def test_teardown_calls_the_pp_purge():
+ from vllm.v1.worker.gpu import cudagraph_utils
+
+ src = inspect.getsource(cudagraph_utils._teardown_profiling_state)
+ assert "drop_deferred_states()" in src
+ assert "recoverssm.reset()" in src
+
+
+def test_context_binding_guard():
+ good = torch.zeros(4)
+ ctx = SimpleNamespace(checkpoints=(good,))
+ gdn_mod._check_context_binding(ctx, [good])
+ with pytest.raises(RuntimeError, match="stale context"):
+ gdn_mod._check_context_binding(ctx, [torch.zeros(4)])
+
+
+def _calls_in(fn) -> list[str]:
+ tree = ast.parse(textwrap.dedent(inspect.getsource(fn)))
+ names = []
+ for node in ast.walk(tree):
+ if isinstance(node, ast.Call) and isinstance(node.func, ast.Attribute):
+ names.append((node.lineno, node.col_offset, node.func.attr))
+ return [n for _, _, n in sorted(names)]
+
+
+def test_runner_wiring_order():
+ """The runner detaches after receive() and before the immediate prefill-
+ chunk postprocess (which would otherwise drop the step), and restores
+ before the deferred postprocess."""
+ from vllm.v1.worker.gpu.model_runner import GPUModelRunner
+
+ calls = _calls_in(GPUModelRunner.sample_tokens)
+ i_recv = calls.index("receive")
+ i_attach = calls.index("attach_deferred_state")
+ i_detach = calls.index("detach_deferred_state")
+ i_post = calls.index("postprocess_state")
+ assert i_recv < i_attach and i_recv < i_detach < i_post and i_attach < i_post
+
+ calls = _calls_in(GPUModelRunner.update_pp_decode_requests)
+ assert calls.index("restore_deferred_state") < calls.index("postprocess_sampled")
diff --git a/vllm/config/vllm.py b/vllm/config/vllm.py
index f215ab3..d07579e 100644
--- a/vllm/config/vllm.py
+++ b/vllm/config/vllm.py
@@ -3272,9 +3272,25 @@ class VllmConfig:
"RecoverSSM with align mode requires VLLM_USE_V2_MODEL_RUNNER=1"
)
if self.parallel_config.pipeline_parallel_size > 1:
- raise ValueError(
- "RecoverSSM currently requires pipeline_parallel_size=1"
- )
+ # Non-last PP ranks get the accepted counts pp_size steps after
+ # their verify. Model Runner V2 defers the commit with a
+ # snapshot of its inputs (RecoverSSMState.detach_step) and
+ # masks freed rows; only the Qwen4Exp GDN/PLE metadata
+ # implement that snapshot (Kimi-K3 KDA does not).
+ if (
+ self.model_config is None
+ or self.model_config.architecture
+ not in _QWEN4EXP_RECOVERSSM_ARCHS
+ ):
+ raise ValueError(
+ "RecoverSSM with pipeline_parallel_size > 1 is supported "
+ "only for Qwen4Exp GDN/PLE"
+ )
+ if not self.use_v2_model_runner:
+ raise ValueError(
+ "RecoverSSM with pipeline_parallel_size > 1 requires "
+ "VLLM_USE_V2_MODEL_RUNNER=1"
+ )
if self.mamba_config.backend != MambaBackendEnum.TRITON:
raise ValueError("RecoverSSM requires --mamba-backend triton")
elif self.cache_config.mamba_cache_mode == "all":
diff --git a/vllm/v1/attention/backends/gdn_recoverssm.py b/vllm/v1/attention/backends/gdn_recoverssm.py
index 8b6349a..e1c12ac 100644
--- a/vllm/v1/attention/backends/gdn_recoverssm.py
+++ b/vllm/v1/attention/backends/gdn_recoverssm.py
@@ -8,7 +8,7 @@ column, num_accepted_tokens is always 1 (the checkpoint and the compacted conv w
already hold the accepted state), and the metadata carries what the post-sampling commit
needs. Mirrors the Kimi-K3 KDA builder's RecoverSSM handling."""
-from dataclasses import dataclass, field, fields
+from dataclasses import dataclass, field, fields, replace
from typing import Any
import torch
@@ -23,6 +23,7 @@ from vllm.v1.attention.backends.gdn_attn import (
from vllm.v1.attention.backends.recoverssm_metadata import (
RecoverSSMMetadata,
RecoverSSMPostprocessMetadata,
+ clone_commit_tensors,
recoverssm_request_indices,
recoverssm_spec_rows,
)
@@ -74,6 +75,32 @@ class GDNRecoverSSMMetadata(GDNAttentionMetadata, RecoverSSMMetadata):
block_size=c.block_size,
)
+ def snapshot_for_deferred_commit(self) -> "GDNRecoverSSMMetadata":
+ c = self.recoverssm_commit
+ if c is None:
+ return self
+ return replace(
+ self,
+ recoverssm_commit=clone_commit_tensors(c, self.num_spec_decodes),
+ )
+
+
+def _check_context_binding(context: Any, checkpoints: list[torch.Tensor]) -> None:
+ """A cached commit context must point at the layers' live KV cache.
+
+ The context stores raw base addresses of every layer's conv, state and
+ replay tensors. One that outlived its cache (the profiling KV cache is
+ torn down and re-allocated before serving) would keep that cache alive
+ and, worse, commit into freed memory. Builders are re-created with the
+ real cache, so a mismatch is a lifecycle bug: fail loudly."""
+ bound = [t.data_ptr() for t in context.checkpoints]
+ live = [t.data_ptr() for t in checkpoints]
+ if bound != live:
+ raise RuntimeError(
+ "GDN RecoverSSM commit context is bound to a different KV cache than "
+ "its layers (stale context after a KV cache re-allocation)"
+ )
+
class GDNRecoverSSMMetadataBuilder(GDNAttentionMetadataBuilder):
# Uniform spec-decode batches replay as FULL graphs: the base builder stages the
@@ -91,6 +118,11 @@ class GDNRecoverSSMMetadataBuilder(GDNAttentionMetadataBuilder):
self._capturing = False
def _get_context(self):
+ if self._context is not None:
+ fc = self.vllm_config.compilation_config.static_forward_context
+ _check_context_binding(
+ self._context, [fc[name].kv_cache[1] for name in self.layer_names]
+ )
if self._context is None:
from vllm.model_executor.layers.mamba.gdn.recoverssm_gdn import (
GDNRecoverSSMCommitContext,
diff --git a/vllm/v1/attention/backends/ple_recoverssm.py b/vllm/v1/attention/backends/ple_recoverssm.py
index 9b43a6f..9e44ad8 100644
--- a/vllm/v1/attention/backends/ple_recoverssm.py
+++ b/vllm/v1/attention/backends/ple_recoverssm.py
@@ -9,7 +9,7 @@ every active decode row takes the spec path, the conv reads at offset 0, and aft
sampling the accepted window is compacted to the front (the shared RecoverSSM compaction
kernel; in align mode it also writes the block-boundary window)."""
-from dataclasses import dataclass, field, fields
+from dataclasses import dataclass, field, fields, replace
from typing import Any
import torch
@@ -25,6 +25,7 @@ from vllm.v1.attention.backend import AttentionCGSupport
from vllm.v1.attention.backends.recoverssm_metadata import (
RecoverSSMMetadata,
RecoverSSMPostprocessMetadata,
+ clone_commit_tensors,
recoverssm_request_indices,
recoverssm_spec_rows,
)
@@ -228,6 +229,15 @@ class PleRecoverSSMMetadata(PleShortConvAttentionMetadata, RecoverSSMMetadata):
block_size=c.block_size,
)
+ def snapshot_for_deferred_commit(self) -> "PleRecoverSSMMetadata":
+ c = self.recoverssm_commit
+ if c is None:
+ return self
+ return replace(
+ self,
+ recoverssm_commit=clone_commit_tensors(c, self.num_spec_decodes),
+ )
+
class PleRecoverSSMMetadataBuilder(PleShortConvAttentionMetadataBuilder):
# Same as the GDN RecoverSSM builder: uniform spec-decode batches replay as FULL
@@ -245,6 +255,17 @@ class PleRecoverSSMMetadataBuilder(PleShortConvAttentionMetadataBuilder):
self._capturing = False
def _get_ctx(self):
+ if self._ctx is not None:
+ fc = self.vllm_config.compilation_config.static_forward_context
+ live = [fc[name].kv_cache[0].data_ptr() for name in self.layer_names]
+ if [s.data_ptr() for s in self._ctx.conv_states] != live:
+ # Same lifecycle guard as the GDN builder: never commit into a
+ # KV cache that was torn down (profiling) and re-allocated.
+ raise RuntimeError(
+ "PLE RecoverSSM commit context is bound to a different KV "
+ "cache than its layers (stale context after a KV cache "
+ "re-allocation)"
+ )
if self._ctx is None:
fc = self.vllm_config.compilation_config.static_forward_context
convs = [fc[name].kv_cache[0] for name in self.layer_names]
diff --git a/vllm/v1/attention/backends/recoverssm_metadata.py b/vllm/v1/attention/backends/recoverssm_metadata.py
index b1339e9..135b3eb 100644
--- a/vllm/v1/attention/backends/recoverssm_metadata.py
+++ b/vllm/v1/attention/backends/recoverssm_metadata.py
@@ -2,7 +2,7 @@
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
import abc
-from dataclasses import dataclass
+from dataclasses import dataclass, replace
import torch
@@ -25,6 +25,25 @@ class RecoverSSMMetadata(abc.ABC):
) -> RecoverSSMPostprocessMetadata | None:
raise NotImplementedError
+ def has_recoverssm_commit(self) -> bool:
+ """Whether this step's metadata carries a commit (any spec row)."""
+ return getattr(self, "recoverssm_commit", None) is not None
+
+ def snapshot_for_deferred_commit(self) -> "RecoverSSMMetadata":
+ """A copy whose commit inputs no longer alias per-step buffers.
+
+ Under pipeline parallelism a non-last rank learns the accepted counts
+ pp_size steps after its verify (the PP sampled-token broadcast), and
+ the steps in between rebuild the builders' persistent, graph-staged
+ buffers (spec state indices, query offsets, the batch-order block
+ table). The deferred commit must read this step's values, so they are
+ cloned here, on the current stream, before the next build overwrites
+ them. Implementations that do not support the PP deferral raise."""
+ raise NotImplementedError(
+ f"{type(self).__name__} does not support the pipeline-parallel "
+ "RecoverSSM commit deferral"
+ )
+
def recoverssm_spec_rows(
is_prefilling_cpu: torch.Tensor,
@@ -62,3 +81,23 @@ def recoverssm_request_indices(
if num_spec_decodes == 0 or int(rows[-1]) == num_spec_decodes - 1:
return None
return rows.to(torch.int32).to(device, non_blocking=True)
+
+
+def clone_commit_tensors(commit, num_spec_decodes: int):
+ """Clone the commit inputs a deferred (PP) commit reads, trimmed to the real
+ spec rows. Works for every RecoverSSM commit dataclass with the fields
+ state_indices, query_start_loc, request_indices, block_table and
+ num_computed_tokens."""
+ n = num_spec_decodes
+
+ def c(t: torch.Tensor | None) -> torch.Tensor | None:
+ return None if t is None else t.clone()
+
+ return replace(
+ commit,
+ state_indices=commit.state_indices[:n].clone(),
+ query_start_loc=commit.query_start_loc[: n + 1].clone(),
+ request_indices=c(commit.request_indices),
+ block_table=c(commit.block_table),
+ num_computed_tokens=c(commit.num_computed_tokens),
+ )
diff --git a/vllm/v1/worker/gpu/cudagraph_utils.py b/vllm/v1/worker/gpu/cudagraph_utils.py
index b28286a..924acee 100644
--- a/vllm/v1/worker/gpu/cudagraph_utils.py
+++ b/vllm/v1/worker/gpu/cudagraph_utils.py
@@ -1119,6 +1119,10 @@ def _teardown_profiling_state(runner: "GPUModelRunner") -> None:
# profiling KV cache.
if (recoverssm := getattr(runner.model_state, "recoverssm", None)) is not None:
recoverssm.reset()
+ # Under PP the non-last ranks keep detached RecoverSSM steps on the pending
+ # PP entries; they reference the profiling KV cache just the same.
+ if (pp_handler := getattr(runner, "pp_handler", None)) is not None:
+ pp_handler.drop_deferred_states()
# Invalidate the align-mode Mamba group metadata cached from the
# profiling KVCacheConfig: the real (e.g. PP-projected) config may
# place Mamba layers into a different group layout, so it must be
diff --git a/vllm/v1/worker/gpu/model_runner.py b/vllm/v1/worker/gpu/model_runner.py
index 246fe4c..1876f4b 100644
--- a/vllm/v1/worker/gpu/model_runner.py
+++ b/vllm/v1/worker/gpu/model_runner.py
@@ -1202,6 +1202,11 @@ class GPUModelRunner(LoRAModelRunnerMixin):
self.req_states.draft_tokens
)
if outputs is not None:
+ # Restore the step's deferred model state (RecoverSSM commits)
+ # so this postprocess commits the step that sampled them.
+ self.model_state.restore_deferred_state(
+ outputs.pop("deferred_state", None)
+ )
self.postprocess_sampled(**outputs)
def add_requests(self, scheduler_output: SchedulerOutput) -> None:
@@ -2306,6 +2311,12 @@ class GPUModelRunner(LoRAModelRunnerMixin):
):
raise RuntimeError("PLE test fault: exc_async_output")
all_decode_next = self.pp_handler.receive(input_batch)
+ # The accepted counts of this step arrive pp_size steps later; keep
+ # the state their postprocess needs with this step's pending entry
+ # (before the immediate prefill-chunk postprocess below).
+ self.pp_handler.attach_deferred_state(
+ self.model_state.detach_deferred_state()
+ )
# Optimistically update num_computed_tokens for entire batch here.
# Will be adjusted for rejections if necessary in update_requests.
self.postprocess_num_computed_tokens(input_batch)
diff --git a/vllm/v1/worker/gpu/model_states/interface.py b/vllm/v1/worker/gpu/model_states/interface.py
index 8ce5121..6f41111 100644
--- a/vllm/v1/worker/gpu/model_states/interface.py
+++ b/vllm/v1/worker/gpu/model_states/interface.py
@@ -163,6 +163,19 @@ class ModelState(ABC):
) -> None:
return None
+ def detach_deferred_state(self) -> Any:
+ """Non-last PP rank, after this step's forward: hand over any state the
+ step's postprocess needs once its sampled tokens arrive (pp_size steps
+ later). The PP handler keeps it with the step's pending entry. None by
+ default (nothing deferred)."""
+ return None
+
+ def restore_deferred_state(self, state: Any) -> None:
+ """Counterpart of `detach_deferred_state`, called right before the
+ deferred postprocess of the step that detached `state`."""
+ if state is not None:
+ raise RuntimeError(f"{type(self).__name__} cannot restore {state!r}")
+
@abstractmethod
def prepare_inputs_embeds(
self,
diff --git a/vllm/v1/worker/gpu/model_states/mamba_hybrid.py b/vllm/v1/worker/gpu/model_states/mamba_hybrid.py
index de0e1ce..b31a0ce 100644
--- a/vllm/v1/worker/gpu/model_states/mamba_hybrid.py
+++ b/vllm/v1/worker/gpu/model_states/mamba_hybrid.py
@@ -359,6 +359,20 @@ class MambaHybridModelState(DefaultModelState):
)
return attn_metadata
+ def detach_deferred_state(self) -> Any:
+ # Non-last PP rank: the RecoverSSM commit of this step has to wait for
+ # the accepted counts of the PP broadcast; snapshot its inputs now.
+ if self.recoverssm is None:
+ return None
+ return self.recoverssm.detach_step()
+
+ def restore_deferred_state(self, state: Any) -> None:
+ if state is None:
+ return
+ if self.recoverssm is None:
+ raise RuntimeError("deferred RecoverSSM state without RecoverSSM")
+ self.recoverssm.restore_step(state)
+
def postprocess_state(
self,
idx_mapping: torch.Tensor,
diff --git a/vllm/v1/worker/gpu/model_states/recoverssm.py b/vllm/v1/worker/gpu/model_states/recoverssm.py
index 3cfbece..2aafa25 100644
--- a/vllm/v1/worker/gpu/model_states/recoverssm.py
+++ b/vllm/v1/worker/gpu/model_states/recoverssm.py
@@ -1,25 +1,77 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+from dataclasses import dataclass
from typing import Any
import torch
+from vllm.logger import init_logger
from vllm.triton_utils import tl, triton
from vllm.v1.attention.backends.recoverssm_metadata import RecoverSSMMetadata
from vllm.v1.worker.utils import AttentionGroup
+logger = init_logger(__name__)
+
+
+@dataclass(frozen=True)
+class DeferredRecoverSSMStep:
+ """One verify step's RecoverSSM commits, detached on a non-last PP rank.
+
+ The rank verified the window at step T but receives the accepted counts
+ only when the PP broadcast of step T is consumed (pp_size steps later).
+ The metadata is a snapshot (`snapshot_for_deferred_commit`), so the steps
+ in between cannot overwrite its commit inputs. It rides on the PP
+ handler's pending entry for step T and is restored right before that
+ entry's postprocess."""
+
+ metadata: tuple[RecoverSSMMetadata, ...]
+
class RecoverSSMState:
"""Coordinates RecoverSSM metadata between attention and postprocessing."""
def __init__(self) -> None:
self._step: tuple[RecoverSSMMetadata, ...] | None = None
+ # True while `_step` is a restored PP-deferred step: its commit must
+ # skip rows whose request was freed since the verify (their state
+ # blocks may already belong to another request).
+ self._step_is_deferred = False
def reset(self) -> None:
"""Drop the recorded step. Its metadata holds the commit contexts, which
reference the KV cache tensors they were built for."""
self._step = None
+ self._step_is_deferred = False
+
+ def detach_step(self) -> DeferredRecoverSSMStep | None:
+ """Take this step's commits for a deferred (non-last PP rank) commit.
+
+ Returns None when the step has nothing to commit (no spec rows)."""
+ step = self._step
+ self._step = None
+ self._step_is_deferred = False
+ if step is None:
+ return None
+ live = tuple(m for m in step if m.has_recoverssm_commit())
+ if not live:
+ return None
+ return DeferredRecoverSSMStep(
+ tuple(m.snapshot_for_deferred_commit() for m in live)
+ )
+
+ def restore_step(self, deferred: DeferredRecoverSSMStep | None) -> None:
+ """Make a detached step the one the next `commit_step` commits."""
+ if self._step is not None:
+ # Every non-last-rank step detaches its record in sample_tokens, so
+ # only a capture-time record (for_capture=False under PIECEWISE)
+ # can be here; it has no sampled tokens and is never committed.
+ logger.warning_once(
+ "RecoverSSM: dropping a recorded step that was never committed "
+ "(expected only after a PIECEWISE capture)"
+ )
+ self._step = None if deferred is None else deferred.metadata
+ self._step_is_deferred = deferred is not None
def record_step(
self,
@@ -49,10 +101,26 @@ class RecoverSSMState:
num_accepted_tokens: torch.Tensor,
) -> None:
step = self._step
+ deferred = self._step_is_deferred
self._step = None
+ self._step_is_deferred = False
if isinstance(num_sampled, int) or step is None:
+ if deferred and step is not None:
+ raise RuntimeError(
+ "RecoverSSM: a deferred PP commit was restored without "
+ "sampled counts"
+ )
return
+ if deferred:
+ # Under PP the consume filter marks freed/non-sampling rows with -1
+ # in idx_mapping. A zero count makes the commit plan skip the row
+ # (commit_len 0): no state, boundary or conv write into blocks the
+ # request no longer owns.
+ num_sampled = torch.where(
+ idx_mapping >= 0, num_sampled, torch.zeros_like(num_sampled)
+ )
+
for metadata in step:
postprocess_meta = metadata.commit_recoverssm_state(num_sampled)
if postprocess_meta is None:
diff --git a/vllm/v1/worker/gpu/pp_utils.py b/vllm/v1/worker/gpu/pp_utils.py
index ab1b991..0a54342 100644
--- a/vllm/v1/worker/gpu/pp_utils.py
+++ b/vllm/v1/worker/gpu/pp_utils.py
@@ -3,7 +3,7 @@
"""Pipeline Parallelism utils for V2 Model Runner."""
from collections import deque
-from dataclasses import dataclass
+from dataclasses import dataclass, replace
import numpy as np
import torch
@@ -32,6 +32,9 @@ class PendingRecv:
# detect requests aborted since then.
gen_at_receive_np: np.ndarray # [num_reqs]
draft_tokens: torch.Tensor | None = None # [num_reqs, num_speculative_steps]
+ # Model-state data the step's postprocess needs once the sampled tokens
+ # arrive (RecoverSSM: the snapshot of the step's deferred commits).
+ deferred_state: object | None = None
def compute_need_sampled_mask(input_batch: InputBatch) -> np.ndarray | None:
@@ -158,8 +161,34 @@ class PPHandler:
num_sampled=slot.num_sampled,
num_rejected=slot.num_rejected,
idx_mapping=idx_mapping,
+ deferred_state=slot.deferred_state,
)
+ def attach_deferred_state(self, state: object | None) -> None:
+ """Attach model-state data to the entry `receive` just reserved for this
+ step. Called right after `receive` in the same step."""
+ assert not self.is_last_rank
+ if state is None:
+ return
+ slot = self.queue[-1]
+ if slot is None:
+ # receive() skipped the broadcast: no row of this step samples, so
+ # no row can have a speculative window to commit.
+ raise RuntimeError(
+ "deferred model state for a step without sampled outputs"
+ )
+ self.queue[-1] = replace(slot, deferred_state=state)
+
+ def drop_deferred_states(self) -> int:
+ """Release every pending entry's deferred state (profiling teardown: the
+ state references the KV cache it was built for). Returns the count."""
+ dropped = 0
+ for i, slot in enumerate(self.queue):
+ if slot is not None and slot.deferred_state is not None:
+ self.queue[i] = replace(slot, deferred_state=None)
+ dropped += 1
+ return dropped
+
def broadcast_drafts(
self, draft_tokens: torch.Tensor, input_batch: InputBatch
) -> None:
--
2.43.0
Reference in New Issue View Git Blame Copy Permalink