From a80e41d690a7f6dde281e0fd52e2dae1e91ab33d Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?J=C3=BCrgen=20Schmied?= Date: Sun, 27 Sep 2026 00:07:04 +0200 Subject: [PATCH 1/3] [Mamba] Share RecoverSSM commit kernels and helpers across models MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit RecoverSSM's commit-plan and conv-window compaction kernels are model independent. Move them from models/kimi_k3 into model_executor/layers/mamba/ops/recoverssm.py (Kimi-K3 keeps its private names as import aliases; kernel source unchanged), add mamba/recoverssm_utils.py (activation predicate, argument validation) and the generic spec-row/request mapping helpers to v1/attention/backends/recoverssm_metadata.py. CacheConfig.use_kda_recoverssm becomes use_recoverssm: the switch now selects RecoverSSM for any model whose recurrent layers implement it, not only KDA. No functional change for Kimi-K3. Signed-off-by: Jürgen Schmied Co-authored-by: Claude --- tests/models/kimi_k3/test_kda.py | 2 +- tests/models/kimi_k3/test_kda_metadata.py | 6 +- tests/test_config.py | 4 +- tests/v1/cudagraph/test_cudagraph_manager.py | 8 +- vllm/config/cache.py | 2 +- vllm/model_executor/layers/mamba/abstract.py | 2 +- .../layers/mamba/ops/recoverssm.py | 196 ++++++++++++++++++ .../layers/mamba/recoverssm_utils.py | 25 +++ vllm/models/kimi_k3/nvidia/kda.py | 6 +- vllm/models/kimi_k3/nvidia/kda_metadata.py | 2 +- vllm/models/kimi_k3/nvidia/model.py | 4 +- vllm/models/kimi_k3/nvidia/ops/recoverssm.py | 178 +--------------- .../attention/backends/recoverssm_metadata.py | 38 ++++ vllm/v1/worker/gpu/cudagraph_utils.py | 2 +- .../worker/gpu/model_states/mamba_hybrid.py | 2 +- 15 files changed, 285 insertions(+), 192 deletions(-) create mode 100644 vllm/model_executor/layers/mamba/ops/recoverssm.py create mode 100644 vllm/model_executor/layers/mamba/recoverssm_utils.py diff --git a/tests/models/kimi_k3/test_kda.py b/tests/models/kimi_k3/test_kda.py index eaf4b69b40..65576c2809 100644 --- a/tests/models/kimi_k3/test_kda.py +++ b/tests/models/kimi_k3/test_kda.py @@ -105,7 +105,7 @@ def test_kda_recoverssm_config_state_layout(): cache_config=SimpleNamespace( mamba_cache_dtype="auto", mamba_ssm_cache_dtype="auto", - use_kda_recoverssm=True, + use_recoverssm=True, ), parallel_config=SimpleNamespace(tensor_parallel_size=1), speculative_config=SimpleNamespace(num_speculative_tokens=2), diff --git a/tests/models/kimi_k3/test_kda_metadata.py b/tests/models/kimi_k3/test_kda_metadata.py index 6268cb9bcc..a5133129f6 100644 --- a/tests/models/kimi_k3/test_kda_metadata.py +++ b/tests/models/kimi_k3/test_kda_metadata.py @@ -119,7 +119,7 @@ def _make_builder( ) vllm_config.cache_config.mamba_cache_mode = mamba_cache_mode vllm_config.cache_config.use_replayssm = use_recoverssm - vllm_config.cache_config.use_kda_recoverssm = use_recoverssm + vllm_config.cache_config.use_recoverssm = use_recoverssm vllm_config.cache_config.prefix_match_unit = prefix_match_unit builder = builder_cls( kv_cache_spec=MambaSpec( @@ -161,7 +161,7 @@ def test_kda_recoverssm_startup_metadata_flow_without_model(monkeypatch): cache_config=SimpleNamespace( mamba_cache_dtype="auto", mamba_ssm_cache_dtype="auto", - use_kda_recoverssm=True, + use_recoverssm=True, ), parallel_config=SimpleNamespace(tensor_parallel_size=1), speculative_config=SimpleNamespace(num_speculative_tokens=2), @@ -194,7 +194,7 @@ def test_kda_recoverssm_startup_metadata_flow_without_model(monkeypatch): ), cache_config=SimpleNamespace( mamba_cache_mode="align", - use_kda_recoverssm=True, + use_recoverssm=True, prefix_match_unit=None, ), parallel_config=SimpleNamespace(decode_context_parallel_size=1), diff --git a/tests/test_config.py b/tests/test_config.py index bbb12e3917..e9540396eb 100644 --- a/tests/test_config.py +++ b/tests/test_config.py @@ -150,7 +150,7 @@ def test_kda_recoverssm_derivation_is_revalidated(): config = SimpleNamespace( cache_config=SimpleNamespace( use_replayssm=True, - use_kda_recoverssm=False, + use_recoverssm=False, mamba_cache_mode="none", ), num_speculative_tokens=3, @@ -169,7 +169,7 @@ def test_kda_recoverssm_derivation_is_revalidated(): VllmConfig.validate_mamba_cached_kernel(config) assert config.cache_config.use_replayssm - assert config.cache_config.use_kda_recoverssm + assert config.cache_config.use_recoverssm config.cache_config.mamba_cache_mode = "align" VllmConfig.validate_mamba_cached_kernel(config) diff --git a/tests/v1/cudagraph/test_cudagraph_manager.py b/tests/v1/cudagraph/test_cudagraph_manager.py index d46476f8d0..411a4f06ec 100644 --- a/tests/v1/cudagraph/test_cudagraph_manager.py +++ b/tests/v1/cudagraph/test_cudagraph_manager.py @@ -259,7 +259,7 @@ def _create_decode_vllm_config( num_speculative_tokens: int = 0, dynamic_spec_schedule: list[tuple[int, int, int]] | None = None, max_num_seqs: int = 8, - use_kda_recoverssm: bool = False, + use_recoverssm: bool = False, ) -> MagicMock: compilation_config = CompilationConfig( cudagraph_mode="FULL_AND_PIECEWISE", @@ -274,7 +274,7 @@ def _create_decode_vllm_config( max_num_seqs=max_num_seqs ) vllm_config.parallel_config = ParallelConfig() - vllm_config.cache_config = SimpleNamespace(use_kda_recoverssm=use_kda_recoverssm) + vllm_config.cache_config = SimpleNamespace(use_recoverssm=use_recoverssm) vllm_config.num_speculative_tokens = num_speculative_tokens if dynamic_spec_schedule is None: vllm_config.speculative_config = None @@ -293,7 +293,7 @@ def _make_spec_decode_manager( num_speculative_tokens: int = 0, dynamic_spec_schedule: list[tuple[int, int, int]] | None = None, max_num_seqs: int = 8, - use_kda_recoverssm: bool = False, + use_recoverssm: bool = False, varlen_decode: bool = False, ) -> gpu_cudagraph_utils.CudaGraphManager: monkeypatch.setattr( @@ -312,7 +312,7 @@ def _make_spec_decode_manager( num_speculative_tokens=num_speculative_tokens, dynamic_spec_schedule=dynamic_spec_schedule, max_num_seqs=max_num_seqs, - use_kda_recoverssm=use_kda_recoverssm, + use_recoverssm=use_recoverssm, ), device=torch.device("cpu"), cudagraph_mode=CUDAGraphMode.FULL_AND_PIECEWISE, diff --git a/vllm/config/cache.py b/vllm/config/cache.py index 3121f49532..390277a27e 100644 --- a/vllm/config/cache.py +++ b/vllm/config/cache.py @@ -213,7 +213,7 @@ class CacheConfig: or FlashInfer mamba backend; standard (non-speculative) decode only. In align mode flushes are most efficient when mamba_block_size is a multiple of replayssm_buffer_len, but this is not required.""" - use_kda_recoverssm: bool = field(default=False, init=False) + use_recoverssm: bool = field(default=False, init=False) """Whether Kimi-K3 KDA uses RecoverSSM speculative decode.""" # Will be set after profiling. diff --git a/vllm/model_executor/layers/mamba/abstract.py b/vllm/model_executor/layers/mamba/abstract.py index fa0eb25758..019cde26ce 100644 --- a/vllm/model_executor/layers/mamba/abstract.py +++ b/vllm/model_executor/layers/mamba/abstract.py @@ -78,7 +78,7 @@ class MambaBase(AttentionLayerBase): # never writes the baseline's per-draft-token state slots. num_speculative_blocks=( 0 - if vllm_config.cache_config.use_kda_recoverssm + if vllm_config.cache_config.use_recoverssm else vllm_config.num_speculative_tokens ), ) diff --git a/vllm/model_executor/layers/mamba/ops/recoverssm.py b/vllm/model_executor/layers/mamba/ops/recoverssm.py new file mode 100644 index 0000000000..073bdd55d4 --- /dev/null +++ b/vllm/model_executor/layers/mamba/ops/recoverssm.py @@ -0,0 +1,196 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Model-independent RecoverSSM Triton kernels shared by every recurrent-state +implementation. + +RecoverSSM verifies a speculative window from a read-only checkpoint and, after +sampling, commits the accepted prefix once. The per-model parts (the KDA and GDN verify +kernels and their state reconstruction) live with their models; the kernels they share +are here: + +- ``prepare_commit_plan_kernel``: per request, the commit length and the destination + state blocks (final and, in align mode, the block-boundary block) from the accepted + counts and the block table; +- ``compact_conv_state_kernel``: moves the accepted suffix of an extended short-conv + window to the front of the destination block (and writes the boundary window in align + mode).""" + +from vllm.triton_utils import tl, triton + + +@triton.heuristics( + { + "HAS_REQUEST_INDICES": lambda args: args["request_indices_ptr"] is not None, + "ALIGN_MODE": lambda args: args["block_table_ptr"] is not None, + } +) +@triton.jit +def prepare_commit_plan_kernel( + num_accepted_ptr, + request_indices_ptr, + state_indices_ptr, + query_start_loc_ptr, + block_table_ptr, + num_computed_ptr, + commit_lens_ptr, + final_state_indices_ptr, + boundary_state_indices_ptr, + boundary_recovery_lens_ptr, + null_block_id, + mamba_block_size, + block_table_width, + stride_num_accepted, + stride_request_indices, + stride_state_indices, + stride_query_start_loc, + stride_block_table_row, + stride_block_table_col, + stride_num_computed, + SPEC_QUERY_LEN: tl.constexpr, + HAS_REQUEST_INDICES: tl.constexpr, + ALIGN_MODE: tl.constexpr, +): + spec_idx = tl.program_id(0) + source_state_idx = tl.load(state_indices_ptr + spec_idx * stride_state_indices).to( + tl.int64 + ) + request_idx = spec_idx + if HAS_REQUEST_INDICES: + request_idx = tl.load( + request_indices_ptr + spec_idx * stride_request_indices + ).to(tl.int64) + num_accepted = tl.load(num_accepted_ptr + request_idx * stride_num_accepted).to( + tl.int32 + ) + bos = tl.load(query_start_loc_ptr + spec_idx * stride_query_start_loc).to(tl.int64) + eos = tl.load(query_start_loc_ptr + (spec_idx + 1) * stride_query_start_loc).to( + tl.int64 + ) + query_len = (eos - bos).to(tl.int32) + commit_len = tl.minimum(tl.maximum(num_accepted, 0), query_len) + commit_len = tl.minimum(commit_len, SPEC_QUERY_LEN) + + final_state_idx = source_state_idx + boundary_state_idx = null_block_id + boundary_recovery_len = 0 + if ALIGN_MODE: + num_computed = tl.load(num_computed_ptr + request_idx * stride_num_computed).to( + tl.int32 + ) + final_num_computed = num_computed + commit_len + final_state_col = tl.minimum( + final_num_computed // mamba_block_size, block_table_width - 1 + ) + final_state_idx = tl.load( + block_table_ptr + + request_idx * stride_block_table_row + + final_state_col * stride_block_table_col + ).to(tl.int64) + next_boundary = (num_computed // mamba_block_size + 1) * mamba_block_size + crosses_boundary = final_num_computed >= next_boundary + boundary_recovery_len = next_boundary - num_computed + boundary_state_idx = tl.load( + block_table_ptr + + request_idx * stride_block_table_row + + (next_boundary // mamba_block_size - 1) * stride_block_table_col, + mask=crosses_boundary, + other=null_block_id, + ).to(tl.int64) + valid = (source_state_idx > null_block_id) & (commit_len > 0) + tl.store(commit_lens_ptr + spec_idx, tl.where(valid, commit_len, 0)) + tl.store( + final_state_indices_ptr + spec_idx, + tl.where(valid, final_state_idx, null_block_id), + ) + tl.store( + boundary_state_indices_ptr + spec_idx, + tl.where(valid, boundary_state_idx, null_block_id), + ) + tl.store( + boundary_recovery_lens_ptr + spec_idx, + tl.where(valid, boundary_recovery_len, 0), + ) + + +@triton.jit +def compact_conv_state_kernel( + conv_state_ref_ptr, + conv_state_base_addrs_ptr, + conv_state_block_strides_ptr, + conv_state_dim_strides_ptr, + conv_state_token_strides_ptr, + state_indices_ptr, + commit_lens_ptr, + final_state_indices_ptr, + boundary_state_indices_ptr, + boundary_recovery_lens_ptr, + null_block_id, + conv_dim, + conv_history_len, + stride_state_indices, + BLOCK_D: tl.constexpr, + BLOCK_HISTORY: tl.constexpr, + ALIGN_MODE: tl.constexpr, +): + pid_d = tl.program_id(0) + pid_b = tl.program_id(1) + pid_l = tl.program_id(2) + source_state_idx = tl.load(state_indices_ptr + pid_b * stride_state_indices).to( + tl.int64 + ) + if source_state_idx <= null_block_id: + return + + commit_len = tl.load(commit_lens_ptr + pid_b) + if commit_len == 0: + return + final_state_idx = tl.load(final_state_indices_ptr + pid_b).to(tl.int64) + boundary_state_idx = tl.load(boundary_state_indices_ptr + pid_b).to(tl.int64) + boundary_recovery_len = tl.load(boundary_recovery_lens_ptr + pid_b) + + if final_state_idx <= null_block_id: + return + + base_addr = tl.load(conv_state_base_addrs_ptr + pid_l) + block_stride = tl.load(conv_state_block_strides_ptr + pid_l) + dim_stride = tl.load(conv_state_dim_strides_ptr + pid_l) + token_stride = tl.load(conv_state_token_strides_ptr + pid_l) + conv_state_ptr = base_addr.to(tl.pointer_type(conv_state_ref_ptr.dtype.element_ty)) + source_ptr = conv_state_ptr + source_state_idx * block_stride + final_ptr = conv_state_ptr + final_state_idx * block_stride + + offs_d = pid_d * BLOCK_D + tl.arange(0, BLOCK_D) + offs_h = tl.arange(0, BLOCK_HISTORY) + mask = (offs_d[:, None] < conv_dim) & (offs_h[None, :] < conv_history_len) + final_values = tl.load( + source_ptr + + offs_d[:, None] * dim_stride + + (commit_len - 1 + offs_h[None, :]) * token_stride, + mask=mask, + ) + if ALIGN_MODE: + boundary_values = tl.load( + source_ptr + + offs_d[:, None] * dim_stride + + (boundary_recovery_len - 1 + offs_h[None, :]) * token_stride, + mask=mask & (boundary_state_idx > null_block_id), + ) + boundary_ptr = conv_state_ptr + boundary_state_idx * block_stride + tl.store( + boundary_ptr + + offs_d[:, None] * dim_stride + + offs_h[None, :] * token_stride, + boundary_values, + mask=mask & (boundary_state_idx > null_block_id), + ) + tl.store( + final_ptr + offs_d[:, None] * dim_stride + offs_h[None, :] * token_stride, + final_values, + mask=mask, + ) + + +__all__ = [ + "compact_conv_state_kernel", + "prepare_commit_plan_kernel", +] diff --git a/vllm/model_executor/layers/mamba/recoverssm_utils.py b/vllm/model_executor/layers/mamba/recoverssm_utils.py new file mode 100644 index 0000000000..1f3998e6cf --- /dev/null +++ b/vllm/model_executor/layers/mamba/recoverssm_utils.py @@ -0,0 +1,25 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Model-independent RecoverSSM Python helpers: the activation predicate shared by +layers, models and backends, and argument validation that raises ``ValueError`` +before any raw-pointer kernel launch.""" + + +def recoverssm_require(cond: bool, msg: str, component: str = "RecoverSSM") -> None: + if not cond: + raise ValueError(f"{component}: {msg}") + + +def uses_recoverssm(cache_config, num_speculative_tokens: int) -> bool: + """True when this model instance runs RecoverSSM for its recurrent states. + + VllmConfig sets ``cache_config.use_recoverssm`` for every model whose recurrent + states support it (Kimi-K3 KDA, Qwen GDN/PLE). Without speculative tokens there is + nothing to recover, so the stock path runs.""" + return ( + bool(getattr(cache_config, "use_recoverssm", False)) + and num_speculative_tokens > 0 + ) + + +__all__ = ["recoverssm_require", "uses_recoverssm"] diff --git a/vllm/models/kimi_k3/nvidia/kda.py b/vllm/models/kimi_k3/nvidia/kda.py index 97b4029084..1abafaa85c 100644 --- a/vllm/models/kimi_k3/nvidia/kda.py +++ b/vllm/models/kimi_k3/nvidia/kda.py @@ -486,7 +486,7 @@ class KimiK3DeltaAttention(GatedDeltaNetAttention): self.cache_config.mamba_cache_dtype, self.cache_config.mamba_ssm_cache_dtype, ) - if self.cache_config.use_kda_recoverssm: + if self.cache_config.use_recoverssm: return MambaStateDtypeCalculator.append_kda_recoverssm_record( base_dtypes, self.model_config.dtype ) @@ -502,7 +502,7 @@ class KimiK3DeltaAttention(GatedDeltaNetAttention): conv_kernel_size=self.conv_size, num_spec=self.num_spec, ) - if self.cache_config.use_kda_recoverssm: + if self.cache_config.use_recoverssm: return MambaStateShapeCalculator.append_kda_recoverssm_record( base_shapes, self.num_heads, @@ -521,7 +521,7 @@ class KimiK3DeltaAttention(GatedDeltaNetAttention): run_gemm_rs_ar: bool = False, ) -> None: super().__init__(config, vllm_config, prefix) - self.use_recoverssm = self.cache_config.use_kda_recoverssm + self.use_recoverssm = self.cache_config.use_recoverssm if self.cache_config.use_replayssm and not self.use_recoverssm: raise ValueError( "Kimi-K3 supports --use-replayssm only with speculative decoding" diff --git a/vllm/models/kimi_k3/nvidia/kda_metadata.py b/vllm/models/kimi_k3/nvidia/kda_metadata.py index 2562171109..fa102be0a2 100644 --- a/vllm/models/kimi_k3/nvidia/kda_metadata.py +++ b/vllm/models/kimi_k3/nvidia/kda_metadata.py @@ -327,7 +327,7 @@ class KimiK3KDAMetadataBuilder(GDNAttentionMetadataBuilder): isinstance(additional_config, dict) and additional_config.get("kda_prefill_backend") == "flashinfer" ) - self.use_recoverssm = vllm_config.cache_config.use_kda_recoverssm + self.use_recoverssm = vllm_config.cache_config.use_recoverssm self.spec_state_slots = 1 if self.use_recoverssm else self.num_spec + 1 self.recoverssm_num_accepted_tokens: torch.Tensor | None = None self.recoverssm_context: KDARecoverSSMCommitContext | None = None diff --git a/vllm/models/kimi_k3/nvidia/model.py b/vllm/models/kimi_k3/nvidia/model.py index c1818a444b..d4ed6d2e73 100644 --- a/vllm/models/kimi_k3/nvidia/model.py +++ b/vllm/models/kimi_k3/nvidia/model.py @@ -1637,7 +1637,7 @@ class KimiLinearForCausalLM( vllm_config.cache_config.mamba_cache_dtype, vllm_config.cache_config.mamba_ssm_cache_dtype, ) - if vllm_config.cache_config.use_kda_recoverssm: + if vllm_config.cache_config.use_recoverssm: dtypes = MambaStateDtypeCalculator.append_kda_recoverssm_record( dtypes, vllm_config.model_config.dtype ) @@ -1662,7 +1662,7 @@ class KimiLinearForCausalLM( conv_kernel_size=hf_config.linear_attn_config["short_conv_kernel_size"], num_spec=num_spec, ) - if vllm_config.cache_config.use_kda_recoverssm: + if vllm_config.cache_config.use_recoverssm: return MambaStateShapeCalculator.append_kda_recoverssm_record( shapes, hf_config.linear_attn_config["num_heads"], diff --git a/vllm/models/kimi_k3/nvidia/ops/recoverssm.py b/vllm/models/kimi_k3/nvidia/ops/recoverssm.py index c603ec7ed1..bca9e52f1a 100644 --- a/vllm/models/kimi_k3/nvidia/ops/recoverssm.py +++ b/vllm/models/kimi_k3/nvidia/ops/recoverssm.py @@ -9,6 +9,12 @@ from typing import Any import torch from vllm.model_executor.layers.mamba.mamba_utils import is_conv_state_dim_first +from vllm.model_executor.layers.mamba.ops.recoverssm import ( + compact_conv_state_kernel as _compact_conv_state_kernel, +) +from vllm.model_executor.layers.mamba.ops.recoverssm import ( + prepare_commit_plan_kernel as _prepare_commit_plan_kernel, +) from vllm.triton_utils import tl, triton from vllm.v1.attention.backends.utils import NULL_BLOCK_ID @@ -223,178 +229,6 @@ def _kda_recoverssm_verify_kernel( ) -@triton.heuristics( - { - "HAS_REQUEST_INDICES": lambda args: args["request_indices_ptr"] is not None, - "ALIGN_MODE": lambda args: args["block_table_ptr"] is not None, - } -) -@triton.jit -def _prepare_commit_plan_kernel( - num_accepted_ptr, - request_indices_ptr, - state_indices_ptr, - query_start_loc_ptr, - block_table_ptr, - num_computed_ptr, - commit_lens_ptr, - final_state_indices_ptr, - boundary_state_indices_ptr, - boundary_recovery_lens_ptr, - null_block_id, - mamba_block_size, - block_table_width, - stride_num_accepted, - stride_request_indices, - stride_state_indices, - stride_query_start_loc, - stride_block_table_row, - stride_block_table_col, - stride_num_computed, - SPEC_QUERY_LEN: tl.constexpr, - HAS_REQUEST_INDICES: tl.constexpr, - ALIGN_MODE: tl.constexpr, -): - spec_idx = tl.program_id(0) - source_state_idx = tl.load(state_indices_ptr + spec_idx * stride_state_indices).to( - tl.int64 - ) - request_idx = spec_idx - if HAS_REQUEST_INDICES: - request_idx = tl.load( - request_indices_ptr + spec_idx * stride_request_indices - ).to(tl.int64) - num_accepted = tl.load(num_accepted_ptr + request_idx * stride_num_accepted).to( - tl.int32 - ) - bos = tl.load(query_start_loc_ptr + spec_idx * stride_query_start_loc).to(tl.int64) - eos = tl.load(query_start_loc_ptr + (spec_idx + 1) * stride_query_start_loc).to( - tl.int64 - ) - query_len = (eos - bos).to(tl.int32) - commit_len = tl.minimum(tl.maximum(num_accepted, 0), query_len) - commit_len = tl.minimum(commit_len, SPEC_QUERY_LEN) - - final_state_idx = source_state_idx - boundary_state_idx = null_block_id - boundary_recovery_len = 0 - if ALIGN_MODE: - num_computed = tl.load(num_computed_ptr + request_idx * stride_num_computed).to( - tl.int32 - ) - final_num_computed = num_computed + commit_len - final_state_col = tl.minimum( - final_num_computed // mamba_block_size, block_table_width - 1 - ) - final_state_idx = tl.load( - block_table_ptr - + request_idx * stride_block_table_row - + final_state_col * stride_block_table_col - ).to(tl.int64) - next_boundary = (num_computed // mamba_block_size + 1) * mamba_block_size - crosses_boundary = final_num_computed >= next_boundary - boundary_recovery_len = next_boundary - num_computed - boundary_state_idx = tl.load( - block_table_ptr - + request_idx * stride_block_table_row - + (next_boundary // mamba_block_size - 1) * stride_block_table_col, - mask=crosses_boundary, - other=null_block_id, - ).to(tl.int64) - valid = (source_state_idx > null_block_id) & (commit_len > 0) - tl.store(commit_lens_ptr + spec_idx, tl.where(valid, commit_len, 0)) - tl.store( - final_state_indices_ptr + spec_idx, - tl.where(valid, final_state_idx, null_block_id), - ) - tl.store( - boundary_state_indices_ptr + spec_idx, - tl.where(valid, boundary_state_idx, null_block_id), - ) - tl.store( - boundary_recovery_lens_ptr + spec_idx, - tl.where(valid, boundary_recovery_len, 0), - ) - - -@triton.jit -def _compact_conv_state_kernel( - conv_state_ref_ptr, - conv_state_base_addrs_ptr, - conv_state_block_strides_ptr, - conv_state_dim_strides_ptr, - conv_state_token_strides_ptr, - state_indices_ptr, - commit_lens_ptr, - final_state_indices_ptr, - boundary_state_indices_ptr, - boundary_recovery_lens_ptr, - null_block_id, - conv_dim, - conv_history_len, - stride_state_indices, - BLOCK_D: tl.constexpr, - BLOCK_HISTORY: tl.constexpr, - ALIGN_MODE: tl.constexpr, -): - pid_d = tl.program_id(0) - pid_b = tl.program_id(1) - pid_l = tl.program_id(2) - source_state_idx = tl.load(state_indices_ptr + pid_b * stride_state_indices).to( - tl.int64 - ) - if source_state_idx <= null_block_id: - return - - commit_len = tl.load(commit_lens_ptr + pid_b) - if commit_len == 0: - return - final_state_idx = tl.load(final_state_indices_ptr + pid_b).to(tl.int64) - boundary_state_idx = tl.load(boundary_state_indices_ptr + pid_b).to(tl.int64) - boundary_recovery_len = tl.load(boundary_recovery_lens_ptr + pid_b) - - if final_state_idx <= null_block_id: - return - - base_addr = tl.load(conv_state_base_addrs_ptr + pid_l) - block_stride = tl.load(conv_state_block_strides_ptr + pid_l) - dim_stride = tl.load(conv_state_dim_strides_ptr + pid_l) - token_stride = tl.load(conv_state_token_strides_ptr + pid_l) - conv_state_ptr = base_addr.to(tl.pointer_type(conv_state_ref_ptr.dtype.element_ty)) - source_ptr = conv_state_ptr + source_state_idx * block_stride - final_ptr = conv_state_ptr + final_state_idx * block_stride - - offs_d = pid_d * BLOCK_D + tl.arange(0, BLOCK_D) - offs_h = tl.arange(0, BLOCK_HISTORY) - mask = (offs_d[:, None] < conv_dim) & (offs_h[None, :] < conv_history_len) - final_values = tl.load( - source_ptr - + offs_d[:, None] * dim_stride - + (commit_len - 1 + offs_h[None, :]) * token_stride, - mask=mask, - ) - if ALIGN_MODE: - boundary_values = tl.load( - source_ptr - + offs_d[:, None] * dim_stride - + (boundary_recovery_len - 1 + offs_h[None, :]) * token_stride, - mask=mask & (boundary_state_idx > null_block_id), - ) - boundary_ptr = conv_state_ptr + boundary_state_idx * block_stride - tl.store( - boundary_ptr - + offs_d[:, None] * dim_stride - + offs_h[None, :] * token_stride, - boundary_values, - mask=mask & (boundary_state_idx > null_block_id), - ) - tl.store( - final_ptr + offs_d[:, None] * dim_stride + offs_h[None, :] * token_stride, - final_values, - mask=mask, - ) - - @triton.jit def _commit_kda_state_kernel( state_ref_ptr, diff --git a/vllm/v1/attention/backends/recoverssm_metadata.py b/vllm/v1/attention/backends/recoverssm_metadata.py index 7e4cca68f8..b1339e9523 100644 --- a/vllm/v1/attention/backends/recoverssm_metadata.py +++ b/vllm/v1/attention/backends/recoverssm_metadata.py @@ -24,3 +24,41 @@ class RecoverSSMMetadata(abc.ABC): self, num_accepted_tokens: torch.Tensor ) -> RecoverSSMPostprocessMetadata | None: raise NotImplementedError + + +def recoverssm_spec_rows( + is_prefilling_cpu: torch.Tensor, + query_start_loc_cpu: torch.Tensor, + num_decode_draft_tokens_cpu: torch.Tensor, + num_reqs: int, +) -> tuple[torch.Tensor, torch.Tensor]: + """Which batch rows take the RecoverSSM spec path, shared by the RecoverSSM metadata + builders (GDN, PLE). + + Every active decode row does (a draft-less step is a window of one token), plus + every row that has drafts. Only the first ``num_reqs`` rows are real; padding rows + never take the spec path. Returns the draft-count vector to hand to the base builder + (1 for spec rows, -1 otherwise, padding included) and the spec rows in batch order, + which is also the order of the base builder's spec state indices.""" + query_lens = query_start_loc_cpu[: num_reqs + 1].diff() + active_decode = (~is_prefilling_cpu[:num_reqs]) & (query_lens > 0) + spec_mask = (num_decode_draft_tokens_cpu[:num_reqs] >= 0) | active_decode + drafts = torch.full_like(num_decode_draft_tokens_cpu, -1) + drafts[:num_reqs] = torch.where(spec_mask, 1, -1).to(drafts.dtype) + return drafts, spec_mask.nonzero().flatten() + + +def recoverssm_request_indices( + rows: torch.Tensor, num_spec_decodes: int, device: torch.device +) -> torch.Tensor | None: + """None when the spec rows are exactly 0..n-1, else their batch rows (the commit + plan's request mapping). Raises if the base builder counted a different number of + spec rows than the classification above.""" + if rows.numel() != num_spec_decodes: + raise ValueError( + f"RecoverSSM: {rows.numel()} spec rows classified, but the builder " + f"produced {num_spec_decodes}" + ) + 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) diff --git a/vllm/v1/worker/gpu/cudagraph_utils.py b/vllm/v1/worker/gpu/cudagraph_utils.py index fa1dc4a04b..de85f8fd1d 100644 --- a/vllm/v1/worker/gpu/cudagraph_utils.py +++ b/vllm/v1/worker/gpu/cudagraph_utils.py @@ -342,7 +342,7 @@ class CudaGraphManager: # recoverSSM cannot capture a dummy query wider than its workspace. if mixed_mode and ( - not self.vllm_config.cache_config.use_kda_recoverssm + not self.vllm_config.cache_config.use_recoverssm or num_tokens <= max_decode_tokens ): # for PIECEWISE graphs there is no limit on requests when replaying diff --git a/vllm/v1/worker/gpu/model_states/mamba_hybrid.py b/vllm/v1/worker/gpu/model_states/mamba_hybrid.py index 2b73b55bac..dbd78e6e14 100644 --- a/vllm/v1/worker/gpu/model_states/mamba_hybrid.py +++ b/vllm/v1/worker/gpu/model_states/mamba_hybrid.py @@ -91,7 +91,7 @@ class MambaHybridModelState(DefaultModelState): # columns and the running state_idx are kept GPU-resident. self._align_mode = self.cache_config.mamba_cache_mode == "align" self.recoverssm = ( - RecoverSSMState() if self.cache_config.use_kda_recoverssm else None + RecoverSSMState() if self.cache_config.use_recoverssm else None ) if self._align_mode: self._mamba_state_idx_gpu = torch.zeros( From 03cae7837fec21c776a0331ffce40d7f202364d2 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?J=C3=BCrgen=20Schmied?= Date: Sun, 27 Sep 2026 00:07:04 +0200 Subject: [PATCH 2/3] [Spec Decode][Qwen4Exp] RecoverSSM for Qwen GDN and the PLE short conv MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit With --use-replayssm and speculative decoding, Qwen4Exp's Gated DeltaNet layers verify the draft window from a read-only checkpoint and keep a per-token replay record [correction(V), normalized key(K), gate(1)] in fp32. After sampling, one commit replays the accepted tokens and writes the recurrent state once (plus the block-boundary state in align mode), instead of snapshotting the state after every draft position. The PLE dilated short conv follows the same protocol (accepted-suffix compaction). num_speculative_blocks = 0, so no per-draft Mamba blocks are reserved. - GDN verify/commit kernels: model_executor/layers/mamba/gdn/recoverssm_gdn.py - metadata/backends: v1/attention/backends/gdn_recoverssm.py, ple_recoverssm.py - the verify reads strided q/k/v views of the conv output and writes in place - builders declare AttentionCGSupport.NEVER: FULL_AND_PIECEWISE falls back to PIECEWISE; the fallback now also counts breakable CUDA graphs (empty splitting_ops, with or without torch.compile) as piecewise attention instead of dropping to NONE - the generic Mamba metadata builder skips the Mamba2 ReplaySSM ring-buffer workspaces when --use-replayssm selects RecoverSSM (the PLE short conv uses that builder) - tests: GDN verify/commit vs the native kernels (fp32 and bf16 state, align boundaries, mixed and ragged batches, extreme gates), PLE vs token-by-token decode, config selection Signed-off-by: Jürgen Schmied Co-authored-by: Claude --- tests/config/test_recoverssm_config.py | 173 +++++ tests/kernels/mamba/test_recoverssm_gdn.py | 601 +++++++++++++++ tests/kernels/mamba/test_recoverssm_ple.py | 153 ++++ vllm/config/compilation.py | 13 +- vllm/config/vllm.py | 29 +- .../layers/mamba/gdn/qwen_gdn_linear_attn.py | 132 +++- .../layers/mamba/gdn/recoverssm_gdn.py | 695 ++++++++++++++++++ vllm/model_executor/models/interfaces.py | 1 + vllm/models/qwen4_exp/nvidia/model.py | 59 +- vllm/models/qwen4_exp/nvidia/mtp.py | 4 +- vllm/models/qwen4_exp/nvidia/ple_layer.py | 8 + vllm/v1/attention/backends/gdn_recoverssm.py | 181 +++++ vllm/v1/attention/backends/mamba_attn.py | 7 +- vllm/v1/attention/backends/ple_recoverssm.py | 326 ++++++++ 14 files changed, 2349 insertions(+), 33 deletions(-) create mode 100644 tests/config/test_recoverssm_config.py create mode 100644 tests/kernels/mamba/test_recoverssm_gdn.py create mode 100644 tests/kernels/mamba/test_recoverssm_ple.py create mode 100644 vllm/model_executor/layers/mamba/gdn/recoverssm_gdn.py create mode 100644 vllm/v1/attention/backends/gdn_recoverssm.py create mode 100644 vllm/v1/attention/backends/ple_recoverssm.py diff --git a/tests/config/test_recoverssm_config.py b/tests/config/test_recoverssm_config.py new file mode 100644 index 0000000000..21791aad63 --- /dev/null +++ b/tests/config/test_recoverssm_config.py @@ -0,0 +1,173 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Configuration-level regression tests for Qwen4Exp GDN/PLE RecoverSSM selection. + +``--use-replayssm`` with speculative decoding selects RecoverSSM for Qwen4Exp, as it +does for Kimi-K3 KDA. The MTP drafter's derived config shares ``cache_config`` and must +keep the choice. These tests run the real validator function against stand-in configs. +""" + +from types import SimpleNamespace +from typing import Any, cast + +import pytest + +from vllm.config.mamba import MambaBackendEnum +from vllm.config.vllm import VllmConfig +from vllm.model_executor.layers.mamba.recoverssm_utils import uses_recoverssm +from vllm.v1.attention.backend import AttentionCGSupport + +TARGET = "Qwen4ExpForConditionalGeneration" +DRAFTER = "Qwen4ExpMTP" + + +def _validator(): + validators = VllmConfig.__pydantic_decorators__.model_validators # type: ignore[attr-defined] + return validators["validate_mamba_cached_kernel"].func + + +def _config( + arch: str, + cache_config=None, + num_spec: int = 3, + use_replayssm: bool = True, + supports_replayssm: bool = True, + **overrides, +): + cfg = SimpleNamespace( + num_speculative_tokens=num_spec, + model_config=SimpleNamespace( + architecture=arch, supports_replayssm=supports_replayssm + ), + cache_config=cache_config + or SimpleNamespace( + use_recoverssm=False, + use_replayssm=use_replayssm, + mamba_cache_mode="align", + ), + mamba_config=SimpleNamespace( + enable_stochastic_rounding=False, backend=MambaBackendEnum.TRITON + ), + use_v2_model_runner=True, + parallel_config=SimpleNamespace(pipeline_parallel_size=1), + kv_transfer_config=None, + ) + for key, value in overrides.items(): + target, _, attr = key.partition("__") + if attr: + setattr(getattr(cfg, target), attr, value) + else: + setattr(cfg, target, value) + return cfg + + +def test_validator_is_registered(): + assert "validate_mamba_cached_kernel" in ( + VllmConfig.__pydantic_decorators__.model_validators # type: ignore[attr-defined] + ) + + +@pytest.mark.parametrize( + "arch", ["Qwen4ExpForConditionalGeneration", "Qwen4ExpForCausalLM"] +) +def test_replayssm_with_spec_selects_recoverssm(arch): + cfg = _config(arch) + _validator()(cfg) + assert cfg.cache_config.use_recoverssm is True + assert uses_recoverssm(cfg.cache_config, 3) + + +def test_derived_drafter_keeps_the_choice_and_runs_the_checks(): + target = _config(TARGET) + _validator()(target) + drafter = _config(DRAFTER, cache_config=target.cache_config) + _validator()(drafter) + assert target.cache_config.use_recoverssm is True + bad = _config( + DRAFTER, + cache_config=target.cache_config, + parallel_config__pipeline_parallel_size=2, + ) + with pytest.raises(ValueError, match="pipeline_parallel_size"): + _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"), + ("cache_config__mamba_cache_mode", "all", "none and align"), + ], +) +def test_runtime_checks_reject(override, value, match): + cfg = _config(TARGET, **{override: value}) + with pytest.raises(ValueError, match=match): + _validator()(cfg) + + +def test_flag_unset_keeps_the_stock_path(): + cfg = _config(TARGET, use_replayssm=False) + _validator()(cfg) + assert cfg.cache_config.use_recoverssm is False + assert not uses_recoverssm(cfg.cache_config, 3) + + +def test_qwen4exp_replayssm_requires_speculative_decoding(): + with pytest.raises(ValueError, match="requires speculative decoding"): + _validator()(_config(TARGET, num_spec=0)) + + +def test_architecture_without_replayssm_support_is_rejected(): + cfg = _config("Qwen3NextForCausalLM", supports_replayssm=False) + with pytest.raises(ValueError, match="not supported for architecture"): + _validator()(cfg) + + +def test_uses_recoverssm_predicate(): + on = SimpleNamespace(use_recoverssm=True) + assert uses_recoverssm(on, 3) + assert not uses_recoverssm(on, 0) + assert not uses_recoverssm(SimpleNamespace(use_recoverssm=False), 3) + assert not uses_recoverssm(SimpleNamespace(), 3) + + +def test_builders_fall_back_to_piecewise(): + from vllm.v1.attention.backends.gdn_recoverssm import GDNRecoverSSMMetadataBuilder + from vllm.v1.attention.backends.ple_recoverssm import PleRecoverSSMMetadataBuilder + + for builder in (GDNRecoverSSMMetadataBuilder, PleRecoverSSMMetadataBuilder): + assert ( + builder.get_cudagraph_support(cast(Any, None), cast(Any, None)) + == AttentionCGSupport.NEVER + ) + + +@pytest.mark.parametrize("compile_mode", ["VLLM_COMPILE", "NONE"]) +@pytest.mark.parametrize(("breakable", "expected"), [("1", "PIECEWISE"), ("0", "NONE")]) +def test_never_support_falls_back_to_piecewise_with_breakable_cudagraphs( + monkeypatch, compile_mode, breakable, expected +): + """A backend without full-graph support must keep piecewise graphs when breakable + CUDA graphs split at attention (splitting_ops is empty then, and torch.compile may + be off).""" + from vllm.config.compilation import ( + CompilationConfig, + CompilationMode, + CUDAGraphMode, + ) + + monkeypatch.setenv("VLLM_USE_BREAKABLE_CUDAGRAPH", breakable) + compilation_config = CompilationConfig( + mode=CompilationMode[compile_mode], + cudagraph_mode=CUDAGraphMode.FULL_AND_PIECEWISE, + splitting_ops=[], + ) + resolved = compilation_config.resolve_cudagraph_mode_and_sizes( + AttentionCGSupport.NEVER, + "GDNRecoverSSMAttentionBackend", + uniform_decode_query_len=4, + ) + assert resolved.name == expected diff --git a/tests/kernels/mamba/test_recoverssm_gdn.py b/tests/kernels/mamba/test_recoverssm_gdn.py new file mode 100644 index 0000000000..2488f2bd5c --- /dev/null +++ b/tests/kernels/mamba/test_recoverssm_gdn.py @@ -0,0 +1,601 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""GDN RecoverSSM: verify against the native fused kernel, commit against the native +per-token states, and the boundary checks (malformed metadata must raise ValueError +before any kernel reads or writes out of bounds).""" + +from typing import Any, cast + +import pytest +import torch + +from vllm.model_executor.layers.mamba.gdn.recoverssm_gdn import ( + GDNRecoverSSMCommitContext, + gdn_recoverssm_verify, +) +from vllm.third_party.flash_linear_attention.ops.fused_sigmoid_gating import ( + fused_sigmoid_gating_delta_rule_update, +) + +pytestmark = pytest.mark.skipif(not torch.cuda.is_available(), reason="needs CUDA") +H, HV, K, V, T = 16, 48, 128, 128, 4 +QLENS = (4, 4, 3) + + +def _inputs(nb: int = 10) -> dict: + torch.manual_seed(0) + dev = "cuda" + tot = sum(QLENS) + qsl = torch.tensor([0, 4, 8, 11], dtype=torch.int32, device=dev) + return dict( + A_log=(torch.rand(HV, device=dev) * 2 - 1).float(), + a=torch.randn(tot, HV, device=dev, dtype=torch.bfloat16), + b=torch.randn(tot, HV, device=dev, dtype=torch.bfloat16), + dt_bias=(torch.randn(HV, device=dev) * 0.5).bfloat16(), + q=torch.randn(1, tot, H, K, device=dev, dtype=torch.bfloat16), + k=torch.randn(1, tot, H, K, device=dev, dtype=torch.bfloat16), + v=torch.randn(1, tot, HV, V, device=dev, dtype=torch.bfloat16), + checkpoint_state=(torch.randn(nb, HV, V, K, device=dev) * 0.05).float(), + replay_cache=torch.zeros(nb, HV, T, V + K + 1, device=dev), + query_start_loc=qsl, + state_indices=torch.tensor([2, 5, 7], dtype=torch.int32, device=dev), + spec_query_len=T, + ) + + +def _verify(inp: dict, **override): + d = {**inp, **override} + return gdn_recoverssm_verify( + d["A_log"], + d["a"], + d["b"], + d["dt_bias"], + d["q"], + d["k"], + d["v"], + checkpoint_state=d["checkpoint_state"], + replay_cache=d["replay_cache"], + query_start_loc=d["query_start_loc"], + state_indices=d["state_indices"], + spec_query_len=d["spec_query_len"], + out=d.get("out"), + ) + + +def _native(inp: dict): + return fused_sigmoid_gating_delta_rule_update( + A_log=inp["A_log"], + a=inp["a"], + b=inp["b"], + dt_bias=inp["dt_bias"], + q=inp["q"], + k=inp["k"], + v=inp["v"], + initial_state=inp["checkpoint_state"].clone(), + inplace_final_state=False, + cu_seqlens=inp["query_start_loc"].long(), + ssm_state_indices=inp["state_indices"].long(), + use_qk_l2norm_in_kernel=True, + ) + + +def _context(inp: dict, ckpt: torch.Tensor) -> GDNRecoverSSMCommitContext: + nb = ckpt.shape[0] + conv = [torch.zeros(nb, 8, 3 + T - 1, device="cuda", dtype=torch.bfloat16)] + return GDNRecoverSSMCommitContext.from_tensors( + conv, [ckpt], [inp["replay_cache"]], spec_query_len=T, max_num_reqs=8 + ) + + +def test_verify_matches_native_and_keeps_the_checkpoint(): + inp = _inputs() + before = inp["checkpoint_state"].clone() + out = _verify(inp) + ref, _ = _native(inp) + torch.testing.assert_close( + out.float(), ref.reshape(out.shape).float(), rtol=0, atol=0 + ) + assert torch.equal(inp["checkpoint_state"], before) + + +@pytest.mark.parametrize("accepted", [1, 2, 3, 4]) +def test_commit_matches_native_per_token_state(accepted): + inp = _inputs() + _verify(inp) + _, states = _native(inp) + ckpt = inp["checkpoint_state"].clone() + ctx = _context(inp, ckpt) + acc = torch.tensor( + [min(accepted, q) for q in QLENS], dtype=torch.int32, device="cuda" + ) + ctx.commit(acc, inp["state_indices"], inp["query_start_loc"]) + for i, blk in enumerate(inp["state_indices"].tolist()): + tok = int(inp["query_start_loc"][i]) + int(acc[i]) - 1 + torch.testing.assert_close(ckpt[blk], states[tok], rtol=1e-5, atol=1e-6) + others = [b for b in range(ckpt.shape[0]) if b not in inp["state_indices"].tolist()] + assert torch.equal(ckpt[others], inp["checkpoint_state"][others]) + + +@pytest.mark.parametrize( + "case", ["short_qsl", "capacity", "hv_lt_h", "a_log", "out", "replay", "device"] +) +def test_verify_rejects_malformed_inputs(case): + inp = _inputs() + over, match = { + "short_qsl": ( + dict(query_start_loc=inp["query_start_loc"][:-1]), + "query metadata", + ), + "capacity": ( + dict( + state_indices=inp["state_indices"][:2], + query_start_loc=inp["query_start_loc"][:3], + ), + "activation capacity", + ), + "hv_lt_h": ( + dict( + v=inp["v"][:, :, :8].contiguous(), + a=inp["a"][:, :8].contiguous(), + b=inp["b"][:, :8].contiguous(), + ), + "positive multiple", + ), + "a_log": (dict(A_log=torch.zeros(HV + 1, device="cuda")), "A_log"), + "out": ( + dict( + out=torch.empty(1, 11, HV, V + 1, device="cuda", dtype=torch.bfloat16) + ), + "output shape", + ), + "replay": ( + dict(replay_cache=torch.zeros(10, HV, T, V + K, device="cuda")), + "replay buffer", + ), + "device": (dict(A_log=inp["A_log"].cpu()), "same device"), + }[case] + with pytest.raises(ValueError, match=match): + _verify(inp, **over) + + +def test_commit_rejects_partial_or_short_metadata(): + inp = _inputs() + ctx = _context(inp, inp["checkpoint_state"].clone()) + acc = torch.ones(3, dtype=torch.int32, device="cuda") + bt = torch.zeros(3, 4, dtype=torch.int32, device="cuda") + with pytest.raises(ValueError, match="align metadata is incomplete"): + ctx.commit( + acc, + inp["state_indices"], + inp["query_start_loc"], + block_table=bt, + mamba_block_size=16, + ) + with pytest.raises(ValueError, match="commit metadata"): + ctx.commit(acc, inp["state_indices"], inp["query_start_loc"][:-1]) + with pytest.raises(ValueError, match="request mapping"): + ctx.commit( + acc, + inp["state_indices"], + inp["query_start_loc"], + request_indices=torch.zeros(2, dtype=torch.int32, device="cuda"), + ) + + +def test_context_rejects_inconsistent_layers(): + inp = _inputs() + ck = inp["checkpoint_state"] + conv = torch.zeros(ck.shape[0], 8, 3 + T - 1, device="cuda", dtype=torch.bfloat16) + with pytest.raises(ValueError, match="at least one layer"): + GDNRecoverSSMCommitContext.from_tensors( + [], [], [], spec_query_len=T, max_num_reqs=8 + ) + with pytest.raises(ValueError, match="differ"): + GDNRecoverSSMCommitContext.from_tensors( + [conv], [ck, ck], [inp["replay_cache"]], spec_query_len=T, max_num_reqs=8 + ) + with pytest.raises(ValueError, match="block count"): + GDNRecoverSSMCommitContext.from_tensors( + [conv[:5]], [ck], [inp["replay_cache"]], spec_query_len=T, max_num_reqs=8 + ) + + +# ------------------------------------------------------------------ review round 2: +# mapping, align, extreme gates +from vllm.v1.attention.backends.recoverssm_metadata import ( # noqa: E402 + recoverssm_request_indices, + recoverssm_spec_rows, +) + +# (A_log range, a scale, b offset): drive g = -exp(A_log) * softplus(a + dt_bias) +# and beta = sigmoid(b) +GATES = { + "typical": (0.0, 1.0, 0.0), + "g_near_0": (-12.0, 1.0, 0.0), + "g_very_negative": (2.5, 4.0, 0.0), + "beta_near_0": (0.0, 1.0, -12.0), + "beta_near_1": (0.0, 1.0, 12.0), +} + + +def _make(qlens, nb=16, seed=0, gate="typical"): + g = torch.Generator(device="cuda").manual_seed(seed) + rnd = lambda *s: torch.randn(*s, device="cuda", generator=g) # noqa: E731 + a_log_shift, ascale, boff = GATES[gate] + tot = sum(qlens) + qsl = torch.tensor( + [0, *torch.tensor(qlens).cumsum(0).tolist()], dtype=torch.int32, device="cuda" + ) + return dict( + A_log=( + torch.rand(HV, device="cuda", generator=g) * 2 - 1 + a_log_shift + ).float(), + a=(rnd(tot, HV) * ascale).bfloat16(), + b=(rnd(tot, HV) + boff).bfloat16(), + dt_bias=(rnd(HV) * 0.5).bfloat16(), + q=rnd(1, tot, H, K).bfloat16(), + k=rnd(1, tot, H, K).bfloat16(), + v=rnd(1, tot, HV, V).bfloat16(), + checkpoint_state=(rnd(nb, HV, V, K) * 0.05).float(), + replay_cache=torch.zeros(nb, HV, T, V + K + 1, device="cuda"), + query_start_loc=qsl, + spec_query_len=T, + ) + + +def test_spec_rows_mixed_batch(): + # rows: 0 spec K=3 | 1 prefill | 2 draft-less decode | 3 spec K=1 | 4 padded + # zero-query | 5,6 beyond num_reqs + is_prefilling = torch.tensor([False, True, False, False, False, False, False]) + qsl = torch.tensor([0, 4, 104, 105, 107, 107], dtype=torch.int32) + drafts_in = torch.tensor( + [3, -1, -1, 1, -1, 2, 0], dtype=torch.int32 + ) # 5, 6: garbage padding + drafts, rows = recoverssm_spec_rows(is_prefilling, qsl, drafts_in, num_reqs=5) + assert rows.tolist() == [0, 2, 3] + assert drafts.tolist() == [1, -1, 1, 1, -1, -1, -1] + idx = recoverssm_request_indices(rows, 3, torch.device("cuda")) + assert idx is not None and idx.tolist() == [0, 2, 3] + assert ( + recoverssm_request_indices(torch.tensor([0, 1, 2]), 3, torch.device("cuda")) + is None + ) + with pytest.raises(ValueError, match="spec rows classified"): + recoverssm_request_indices(rows, 2, torch.device("cuda")) + + +def test_commit_maps_noncontiguous_requests_in_align_mode(): + """Spec rows 0, 2, 3 of a 5-request batch; accepted counts, block table and + num_computed are indexed by request row. Every spec request must commit its own + state into its own block.""" + qlens = [4, 1, 2] # the spec rows' windows (K=3, draft-less, K=1) + inp = _make(qlens, nb=40) + bs = 8 + bt = torch.arange(1, 1 + 5 * 6, dtype=torch.int32, device="cuda").reshape( + 5, 6 + ) # request row -> blocks + nc = torch.tensor( + [5, 0, 16, 9, 0], dtype=torch.int32, device="cuda" + ) # per request row + req_rows = torch.tensor([0, 2, 3], dtype=torch.int32, device="cuda") + src = torch.stack([bt[r, int(nc[r]) // bs] for r in req_rows.tolist()]).contiguous() + inp["state_indices"] = src + _verify(inp) + _, states = _native(inp) + ckpt = inp["checkpoint_state"].clone() + ctx = _context(inp, ckpt) + acc = torch.tensor( + [3, 9, 1, 2, 9], dtype=torch.int32, device="cuda" + ) # rows 1, 4 must be ignored + ctx.commit( + acc, + src, + inp["query_start_loc"], + request_indices=req_rows, + block_table=bt, + num_computed_tokens=nc, + mamba_block_size=bs, + ) + for i, r in enumerate(req_rows.tolist()): + n = min(int(acc[r]), qlens[i]) + c0 = int(nc[r]) + final_blk = int(bt[r, (c0 + n) // bs]) + tok = int(inp["query_start_loc"][i]) + n - 1 + torch.testing.assert_close(ckpt[final_blk], states[tok], rtol=1e-5, atol=1e-6) + + +# (num_computed, accepted, block size): boundary after 1 / after 2 / ending exactly on +# it / crossing and continuing / accepted = 1 without crossing / the full window +# crossing after 3 +ALIGN_CASES = [ + (3, 2, 4), + (2, 3, 4), + (0, 4, 4), + (6, 4, 4), + (1, 1, 4), + (5, 4, 8), + (7, 4, 8), + (4, 4, 8), +] + + +@pytest.mark.parametrize("nc,accepted,bs", ALIGN_CASES) +def test_align_boundary_and_final_state_match_native(nc, accepted, bs): + inp = _make([T], nb=12) + bt = torch.arange(2, 8, dtype=torch.int32, device="cuda").reshape(1, 6) + src = bt[0, nc // bs].reshape(1).contiguous() + inp["state_indices"] = src + _verify(inp) + _, states = _native(inp) + ckpt = inp["checkpoint_state"].clone() + ctx = _context(inp, ckpt) + ctx.commit( + torch.tensor([accepted], dtype=torch.int32, device="cuda"), + src, + inp["query_start_loc"], + block_table=bt, + num_computed_tokens=torch.tensor([nc], dtype=torch.int32, device="cuda"), + mamba_block_size=bs, + ) + next_boundary = (nc // bs + 1) * bs + final_blk = int(bt[0, (nc + accepted) // bs]) + if ( + nc + accepted >= next_boundary + ): # crosses: the boundary state goes into the source block + torch.testing.assert_close( + ckpt[int(src)], states[next_boundary - nc - 1], rtol=1e-5, atol=1e-6 + ) + torch.testing.assert_close( + ckpt[final_blk], states[accepted - 1], rtol=1e-5, atol=1e-6 + ) + + +@pytest.mark.parametrize("gate", list(GATES)) +@pytest.mark.parametrize("seed", [0, 1, 2]) +def test_commit_matches_native_under_extreme_gates(gate, seed): + inp = _make(list(QLENS), nb=10, seed=seed, gate=gate) + inp["state_indices"] = torch.tensor([2, 5, 7], dtype=torch.int32, device="cuda") + _verify(inp) + _, states = _native(inp) + for accepted in (1, 2, 3, 4): + ckpt = inp["checkpoint_state"].clone() + ctx = _context(inp, ckpt) + acc = torch.tensor( + [min(accepted, q) for q in QLENS], dtype=torch.int32, device="cuda" + ) + ctx.commit(acc, inp["state_indices"], inp["query_start_loc"]) + for i, blk in enumerate(inp["state_indices"].tolist()): + tok = int(inp["query_start_loc"][i]) + int(acc[i]) - 1 + torch.testing.assert_close(ckpt[blk], states[tok], rtol=1e-5, atol=1e-6) + + +@pytest.mark.parametrize("batch", [1, 2, 4]) +def test_verify_on_strided_views_writes_in_place(batch): + """The layer passes views of the packed conv output (token stride = the packed + width) and the layer's output buffer; result and replay record must equal the + contiguous-copy path.""" + torch.manual_seed(batch) + H, HV, K, V, SQ, NB = 16, 48, 128, 128, 4, 6 + T = batch * SQ + qd, vd = H * K, HV * V + packed = torch.randn(T, 2 * qd + 2 * vd, device="cuda", dtype=torch.bfloat16) + mixed = packed[:, : 2 * qd + vd] + a = torch.randn(T, HV, device="cuda", dtype=torch.bfloat16) + b = torch.randn(T, HV, device="cuda", dtype=torch.bfloat16) + A_log = torch.randn(HV, device="cuda") + dt_bias = torch.randn(HV, device="cuda") + checkpoint = torch.randn(NB, HV, V, K, device="cuda") * 0.05 + qsl = torch.arange(0, T + 1, SQ, device="cuda", dtype=torch.int32) + si = torch.arange(1, batch + 1, device="cuda", dtype=torch.int32) + + def run(q, k, v, out=None): + replay = torch.zeros(NB, HV, SQ, V + K + 1, device="cuda") + o = gdn_recoverssm_verify( + A_log, + a, + b, + dt_bias, + q, + k, + v, + checkpoint_state=checkpoint, + replay_cache=replay, + query_start_loc=qsl, + state_indices=si, + spec_query_len=SQ, + out=out, + ) + return o, replay + + q, k, v = ( + x.contiguous().view(1, T, -1, d) + for x, d in zip(torch.split(mixed, [qd, qd, vd], dim=-1), (K, K, V)) + ) + want_out, want_replay = run(q, k, v) + + core = torch.zeros(T + 3, HV, V, device="cuda", dtype=torch.bfloat16) + views = ( + mixed[:, :qd].view(1, T, -1, K), + mixed[:, qd : 2 * qd].view(1, T, -1, K), + mixed[:, 2 * qd :].view(1, T, -1, V), + ) + got_out, got_replay = run(*views, out=core[:T].unsqueeze(0)) + assert got_out.data_ptr() == core.data_ptr() + assert torch.equal(core[:T].unsqueeze(0), want_out) + assert torch.equal(got_replay, want_replay) + + +def test_verify_on_strided_views_with_variable_query_lengths(): + """Token-strided views with ragged windows (query lengths 4, 1, 3), as MTP + scheduling produces them.""" + torch.manual_seed(7) + H, HV, K, V, SQ, NB = 16, 48, 128, 128, 4, 6 + qsl_list = [0, 4, 5, 8] + T = qsl_list[-1] + qd, vd = H * K, HV * V + packed = torch.randn(T, 2 * qd + 2 * vd, device="cuda", dtype=torch.bfloat16) + mixed = packed[:, : 2 * qd + vd] + a = torch.randn(T, HV, device="cuda", dtype=torch.bfloat16) + b = torch.randn(T, HV, device="cuda", dtype=torch.bfloat16) + A_log = torch.randn(HV, device="cuda") + dt_bias = torch.randn(HV, device="cuda") + checkpoint = torch.randn(NB, HV, V, K, device="cuda") * 0.05 + qsl = torch.tensor(qsl_list, device="cuda", dtype=torch.int32) + si = torch.tensor([1, 3, 5], device="cuda", dtype=torch.int32) + + def run(q, k, v, out=None): + replay = torch.zeros(NB, HV, SQ, V + K + 1, device="cuda") + o = gdn_recoverssm_verify( + A_log, + a, + b, + dt_bias, + q, + k, + v, + checkpoint_state=checkpoint, + replay_cache=replay, + query_start_loc=qsl, + state_indices=si, + spec_query_len=SQ, + out=out, + ) + return o, replay + + q, k, v = ( + x.contiguous().view(1, T, -1, d) + for x, d in zip(torch.split(mixed, [qd, qd, vd], dim=-1), (K, K, V)) + ) + want_out, want_replay = run(q, k, v) + core = torch.zeros(T + 2, HV, V, device="cuda", dtype=torch.bfloat16) + views = ( + mixed[:, :qd].view(1, T, -1, K), + mixed[:, qd : 2 * qd].view(1, T, -1, K), + mixed[:, 2 * qd :].view(1, T, -1, V), + ) + _, got_replay = run(*views, out=core[:T].unsqueeze(0)) + assert torch.equal(core[:T].unsqueeze(0), want_out) + assert torch.equal(got_replay, want_replay) + + +def test_layer_verify_qkv_views_and_fallback(): + """The layer's guard: views for a [tokens, q|k|v] conv output with a contiguous + last dimension, the contiguous-copy fallback for any other layout; same values + either way.""" + from types import SimpleNamespace + + from vllm.model_executor.layers.mamba.gdn.qwen_gdn_linear_attn import ( + QwenGatedDeltaNetAttention as Layer, + ) + + H, HV, K, V, T = 4, 8, 16, 16, 5 + layer = SimpleNamespace( + key_dim=H * K, value_dim=HV * V, tp_size=1, head_k_dim=K, head_v_dim=V + ) + layer.rearrange_mixed_qkv = lambda x: Layer.rearrange_mixed_qkv(cast(Any, layer), x) + width = 2 * H * K + HV * V + packed = torch.randn(T, width + HV * V, dtype=torch.bfloat16) # q|k|v|z, like qkvz + strided = packed[:, :width] + q, k, v, is_view = Layer._recoverssm_verify_qkv(cast(Any, layer), strided) + assert is_view + assert q.data_ptr() == packed.data_ptr() + rq, rk, rv = Layer.rearrange_mixed_qkv(cast(Any, layer), strided) + assert torch.equal(q, rq) and torch.equal(k, rk) and torch.equal(v, rv) + + # Column-major storage: last dimension not contiguous -> copies, same values. + col_major = strided.t().contiguous().t() + assert col_major.stride(-1) != 1 + q2, k2, v2, is_view2 = Layer._recoverssm_verify_qkv(cast(Any, layer), col_major) + assert not is_view2 + assert torch.equal(q2, rq) and torch.equal(k2, rk) and torch.equal(v2, rv) + + +# ------------------------------------------------------------- BF16 recurrent state +# RecoverSSM upcasts the BF16 checkpoint, carries the accepted window in FP32 and rounds +# once when it writes the committed state. The reference is the FP32 trajectory from the +# same (already BF16-rounded) checkpoint; the committed state must equal its BF16 +# rounding up to one BF16 ulp, and be no further from it than the native path is. +BF16_ULP = 2**-7 + + +def _bf16_case(qlens, nb=12, seed=0): + inp = _make(qlens, nb=nb, seed=seed) + inp["checkpoint_state"] = inp["checkpoint_state"].bfloat16() + ref = dict(inp) + ref["checkpoint_state"] = inp["checkpoint_state"].float() + return inp, ref + + +def test_verify_bf16_state_matches_native(): + inp, _ = _bf16_case(list(QLENS)) + inp["state_indices"] = torch.tensor([2, 5, 7], dtype=torch.int32, device="cuda") + before = inp["checkpoint_state"].clone() + out = _verify(inp) + ref, _ = _native(inp) + torch.testing.assert_close( + out.float(), ref.reshape(out.shape).float(), rtol=0, atol=0 + ) + assert torch.equal(inp["checkpoint_state"], before) + + +@pytest.mark.parametrize("accepted", [1, 2, 3, 4]) +@pytest.mark.parametrize("seed", [0, 1]) +def test_commit_bf16_state_rounds_once(accepted, seed): + inp, ref_inp = _bf16_case(list(QLENS), seed=seed) + idx = torch.tensor([2, 5, 7], dtype=torch.int32, device="cuda") + inp["state_indices"] = ref_inp["state_indices"] = idx + _verify(inp) + _, native_states = _native(inp) + _, ref_states = _native(ref_inp) + ckpt = inp["checkpoint_state"].clone() + _context(inp, ckpt).commit( + torch.tensor( + [min(accepted, q) for q in QLENS], dtype=torch.int32, device="cuda" + ), + idx, + inp["query_start_loc"], + ) + for i, blk in enumerate(idx.tolist()): + tok = int(inp["query_start_loc"][i]) + min(accepted, QLENS[i]) - 1 + ref = ref_states[tok].float() + got = ckpt[blk].float() + torch.testing.assert_close( + got, ref.bfloat16().float(), rtol=BF16_ULP, atol=1e-5 + ) + err_rssm = (got - ref).abs().mean() + err_native = (native_states[tok].bfloat16().float() - ref).abs().mean() + assert err_rssm <= err_native * 1.05 + 1e-7, (err_rssm, err_native) + + +@pytest.mark.parametrize("nc,accepted,bs", ALIGN_CASES) +def test_align_boundary_bf16_state(nc, accepted, bs): + inp, ref_inp = _bf16_case([T]) + bt = torch.arange(2, 8, dtype=torch.int32, device="cuda").reshape(1, 6) + src = bt[0, nc // bs].reshape(1).contiguous() + inp["state_indices"] = ref_inp["state_indices"] = src + _verify(inp) + _, ref_states = _native(ref_inp) + ckpt = inp["checkpoint_state"].clone() + _context(inp, ckpt).commit( + torch.tensor([accepted], dtype=torch.int32, device="cuda"), + src, + inp["query_start_loc"], + block_table=bt, + num_computed_tokens=torch.tensor([nc], dtype=torch.int32, device="cuda"), + mamba_block_size=bs, + ) + next_boundary = (nc // bs + 1) * bs + final_blk = int(bt[0, (nc + accepted) // bs]) + if nc + accepted >= next_boundary: + torch.testing.assert_close( + ckpt[int(src)].float(), + ref_states[next_boundary - nc - 1].float().bfloat16().float(), + rtol=BF16_ULP, + atol=1e-5, + ) + torch.testing.assert_close( + ckpt[final_blk].float(), + ref_states[accepted - 1].float().bfloat16().float(), + rtol=BF16_ULP, + atol=1e-5, + ) diff --git a/tests/kernels/mamba/test_recoverssm_ple.py b/tests/kernels/mamba/test_recoverssm_ple.py new file mode 100644 index 0000000000..ec11ad3405 --- /dev/null +++ b/tests/kernels/mamba/test_recoverssm_ple.py @@ -0,0 +1,153 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""PLE short-conv RecoverSSM against the sequential reference. + +Reference: the stock decode path fed the accepted tokens one at a time. +RecoverSSM: one spec-mode conv over the whole window from the checkpoint +(offset 0), then ``_PleConvCommit`` compacts the accepted suffix to the front. +The conv history is pure data movement, so the states must be bit-identical, +and the outputs of the accepted tokens must match as well. +""" + +import pytest +import torch + +from vllm.model_executor.layers.mamba.mamba_utils import is_conv_state_dim_first +from vllm.models.qwen4_exp.nvidia.ops.ple import ple_conv +from vllm.v1.attention.backends.ple_recoverssm import _PleConvCommit + +C, KSIZE, DIL, SQ, NB = 512, 4, 3, 4, 8 +STATE_LEN = (KSIZE - 1) * DIL +WIDTH = STATE_LEN + SQ - 1 +DEV = "cuda" + + +def _state(): + if is_conv_state_dim_first(): + raw = torch.zeros(NB, C, WIDTH, dtype=torch.bfloat16, device=DEV) + return raw, raw + raw = torch.zeros(NB, WIDTH, C, dtype=torch.bfloat16, device=DEV) + return raw, raw.transpose(-1, -2) + + +def _i32(xs): + return torch.tensor(xs, dtype=torch.int32, device=DEV) + + +def _setup(seed, blocks): + g = torch.Generator(device=DEV).manual_seed(seed) + w = torch.randn(C, KSIZE, generator=g, device=DEV).to(torch.bfloat16) + x = torch.randn(len(blocks) * SQ, C, generator=g, device=DEV).to(torch.bfloat16) + outer = torch.randn_like(x) + raw, view = _state() + for b in blocks: + view[b, :, :STATE_LEN] = torch.randn(C, STATE_LEN, generator=g, device=DEV).to( + torch.bfloat16 + ) + return w, x, outer, raw, view + + +def _sequential(w, x, outer, view, block, first, count): + """Stock decode, one token at a time, on a copy of the checkpoint block.""" + ref = view.clone() + res = torch.zeros(count, C, dtype=torch.bfloat16, device=DEV) + for i in range(count): + t = first + i + ple_conv( + inputs=x[t : t + 1], + residual=res[i : i + 1], + conv_state=ref, + conv_weights=w, + state_indices=_i32([block]), + outer_residual=outer[t : t + 1], + mode="decode", + dilation=DIL, + ) + return ref[block, :, :STATE_LEN].clone(), res + + +def _spec(w, x, outer, view, blocks): + res = torch.zeros(len(blocks) * SQ, C, dtype=torch.bfloat16, device=DEV) + ple_conv( + inputs=x, + residual=res, + conv_state=view, + conv_weights=w, + state_indices=_i32(blocks), + outer_residual=outer, + mode="spec", + dilation=DIL, + query_start_loc=_i32(list(range(0, len(blocks) * SQ + 1, SQ))), + num_accepted_tokens=_i32([1] * len(blocks)), + spec_query_len=SQ, + ) + return res + + +@pytest.mark.parametrize("accepted", [1, 2, SQ]) +@pytest.mark.parametrize("seed", [0, 1]) +def test_ple_commit_matches_sequential(accepted, seed): + block = 3 + w, x, outer, raw, view = _setup(seed, [block]) + want_state, want_out = _sequential(w, x, outer, view, block, 0, accepted) + history = torch.cat([view[block, :, :STATE_LEN], x[:accepted].T], dim=1) + + got_out = _spec(w, x, outer, view, [block]) + ctx = _PleConvCommit([raw], spec_query_len=SQ, max_num_reqs=NB) + ctx.commit(_i32([accepted]), _i32([block]), _i32([0, SQ])) + + got_state = view[block, :, :STATE_LEN] + assert torch.equal(got_state, want_state) + assert torch.equal(got_state, history[:, -STATE_LEN:]) + assert torch.equal(got_out[:accepted], want_out) + + +@pytest.mark.parametrize( + ("num_computed", "accepted", "crosses"), + [(6, 3, True), (6, 2, True), (1, 2, False), (7, SQ, True)], +) +def test_ple_commit_align_boundary(num_computed, accepted, crosses): + block_size, src, nxt = 8, 2, 5 + w, x, outer, raw, view = _setup(7, [src]) + boundary_len = block_size - num_computed + want_final, _ = _sequential(w, x, outer, view, src, 0, accepted) + want_boundary, _ = _sequential(w, x, outer, view, src, 0, min(boundary_len, SQ)) + + _spec(w, x, outer, view, [src]) + ctx = _PleConvCommit([raw], spec_query_len=SQ, max_num_reqs=NB) + ctx.commit( + _i32([accepted]), + _i32([src]), + _i32([0, SQ]), + block_table=_i32([[src, nxt]]), + num_computed_tokens=_i32([num_computed]), + mamba_block_size=block_size, + ) + + final_block = nxt if num_computed + accepted >= block_size else src + assert torch.equal(view[final_block, :, :STATE_LEN], want_final) + if crosses: + assert torch.equal(view[src, :, :STATE_LEN], want_boundary) + + +def test_ple_commit_non_contiguous_request_rows(): + blocks, accepted = [1, 6], [3, 1] + w, x, outer, raw, view = _setup(11, blocks) + want = [ + _sequential(w, x, outer, view, b, i * SQ, a)[0] + for i, (b, a) in enumerate(zip(blocks, accepted)) + ] + + _spec(w, x, outer, view, blocks) + # Spec rows 0 and 1 belong to requests 1 and 3 of a four-request batch. + counts = _i32([9, accepted[0], 9, accepted[1]]) + ctx = _PleConvCommit([raw], spec_query_len=SQ, max_num_reqs=NB) + ctx.commit( + counts, + _i32(blocks), + _i32([0, SQ, 2 * SQ]), + request_indices=_i32([1, 3]), + ) + + for b, ref in zip(blocks, want): + assert torch.equal(view[b, :, :STATE_LEN], ref) diff --git a/vllm/config/compilation.py b/vllm/config/compilation.py index eefcf1f49e..3a492843fe 100644 --- a/vllm/config/compilation.py +++ b/vllm/config/compilation.py @@ -1424,10 +1424,15 @@ class CompilationConfig: f"with {min_cg_attn_backend} backend (support: " f"{min_cg_support})" ) - if self.mode == CompilationMode.VLLM_COMPILE and ( - self.splitting_ops_contain_attention() - or self.use_inductor_graph_partition - ): + if ( + self.mode == CompilationMode.VLLM_COMPILE + and ( + self.splitting_ops_contain_attention() + or self.use_inductor_graph_partition + ) + ) or envs.VLLM_USE_BREAKABLE_CUDAGRAPH: + # Breakable CUDA graphs break at the attention kernels at runtime, + # with or without torch.compile. msg += ( "; setting cudagraph_mode=PIECEWISE because " "attention is compiled piecewise" diff --git a/vllm/config/vllm.py b/vllm/config/vllm.py index c50edf4363..494f7f3938 100644 --- a/vllm/config/vllm.py +++ b/vllm/config/vllm.py @@ -70,6 +70,13 @@ else: logger = init_logger(__name__) +# Architectures whose GDN (and PLE) layers run RecoverSSM under --use-replayssm. +_QWEN4EXP_RECOVERSSM_ARCHS = ( + "Qwen4ExpForCausalLM", + "Qwen4ExpForConditionalGeneration", + "Qwen4ExpMTP", +) + # TODO(rocm): These models are either unsupported by MRV2 or slower with # MRV2 on AMD GPUs. ROCM_DEFAULT_MRV1_ARCHITECTURES = frozenset( @@ -3420,21 +3427,35 @@ class VllmConfig: @model_validator(mode="after") def validate_mamba_cached_kernel(self) -> "VllmConfig": if not self.cache_config.use_replayssm: - self.cache_config.use_kda_recoverssm = False + self.cache_config.use_recoverssm = False return self - self.cache_config.use_kda_recoverssm = self.num_speculative_tokens > 0 + self.cache_config.use_recoverssm = self.num_speculative_tokens > 0 if self.model_config is not None and not self.model_config.supports_replayssm: raise ValueError( "--use-replayssm is not supported for architecture " f"{self.model_config.architecture!r}" ) - if self.cache_config.use_kda_recoverssm: + if ( + self.model_config is not None + and self.model_config.architecture in _QWEN4EXP_RECOVERSSM_ARCHS + and not self.cache_config.use_recoverssm + ): + # Qwen4Exp's GDN and PLE layers have no plain ReplaySSM path; with them the + # flag selects RecoverSSM, which needs speculative decoding. + raise ValueError( + "--use-replayssm on Qwen4Exp selects GDN RecoverSSM and requires " + "speculative decoding (num_speculative_tokens > 0)" + ) + if self.cache_config.use_recoverssm: if self.model_config is not None and self.model_config.architecture not in ( "KimiLinearForCausalLM", "KimiK3ForConditionalGeneration", + *_QWEN4EXP_RECOVERSSM_ARCHS, ): - raise ValueError("RecoverSSM is only supported for Kimi-K3 KDA") + raise ValueError( + "RecoverSSM is only supported for Kimi-K3 KDA and Qwen4Exp GDN" + ) if self.mamba_config.enable_stochastic_rounding: raise ValueError( "RecoverSSM supports bfloat16/float32 " diff --git a/vllm/model_executor/layers/mamba/gdn/qwen_gdn_linear_attn.py b/vllm/model_executor/layers/mamba/gdn/qwen_gdn_linear_attn.py index 9d5bcc89da..0f22f0c911 100644 --- a/vllm/model_executor/layers/mamba/gdn/qwen_gdn_linear_attn.py +++ b/vllm/model_executor/layers/mamba/gdn/qwen_gdn_linear_attn.py @@ -370,10 +370,15 @@ class ChunkGatedDeltaRule(CustomOp): @PluggableLayer.register("qwen_gated_delta_net_attention") class QwenGatedDeltaNetAttention(GatedDeltaNetAttention): - def get_state_shape( - self, - ) -> tuple[tuple[int, ...], tuple[int, ...], tuple[int, ...], tuple[int, ...]]: - return MambaStateShapeCalculator.gated_delta_net_state_shape( + def _uses_gdn_recoverssm(self) -> bool: + """One predicate for the replay state's shape, dtype and the attention + backend.""" + from vllm.model_executor.layers.mamba.recoverssm_utils import uses_recoverssm + + return uses_recoverssm(self.cache_config, self.num_spec) + + def get_state_shape(self) -> tuple[tuple[int, ...], ...]: + shapes = MambaStateShapeCalculator.gated_delta_net_state_shape( self.tp_size, self.num_k_heads, self.num_v_heads, @@ -382,6 +387,32 @@ class QwenGatedDeltaNetAttention(GatedDeltaNetAttention): self.conv_kernel_size, self.num_spec, ) + if self._uses_gdn_recoverssm(): + # per-token replay record [HV, spec_query_len, V + K + 1] fp32 + shapes = ( + *shapes, + ( + self.num_v_heads // self.tp_size, + self.num_spec + 1, + self.head_v_dim + self.head_k_dim + 1, + ), + ) + return shapes + + def get_state_dtype(self) -> tuple[torch.dtype, ...]: + dtypes = super().get_state_dtype() + if self._uses_gdn_recoverssm(): + dtypes = (*dtypes, torch.float32) + return dtypes + + def get_attn_backend(self): + if self._uses_gdn_recoverssm(): + from vllm.v1.attention.backends.gdn_recoverssm import ( + GDNRecoverSSMAttentionBackend, + ) + + return GDNRecoverSSMAttentionBackend + return super().get_attn_backend() def __init__( self, @@ -530,6 +561,17 @@ class QwenGatedDeltaNetAttention(GatedDeltaNetAttention): self.gdn_decode_kernel = "XPU" self.enable_fused_gdn_decode = self.gdn_decode_kernel == "cuda" + # RecoverSSM verify replaces the per-draft snapshot path (Triton decode + # only) + self.use_gdn_recoverssm = self._uses_gdn_recoverssm() + if self.use_gdn_recoverssm: + self.gdn_decode_kernel = "triton" + self.enable_fused_gdn_decode = False + self.enable_packed_recurrent_decode = False + logger.info_once( + "GDN RecoverSSM speculative verify active (spec_query_len %d)", + self.num_spec + 1, + ) logger.info_once("GDN decode kernel: %s", self.gdn_decode_kernel) compilation_config = get_current_vllm_config().compilation_config @@ -540,7 +582,7 @@ class QwenGatedDeltaNetAttention(GatedDeltaNetAttention): def _fused_gdn_decode_unsupported_reason( self, vllm_config: VllmConfig ) -> str | None: - conv_state_dtype, recurrent_state_dtype = self.get_state_dtype() + conv_state_dtype, recurrent_state_dtype = self.get_state_dtype()[:2] if ( self.gqa_interleaved_layout or self.head_k_dim != 128 @@ -800,6 +842,25 @@ class QwenGatedDeltaNetAttention(GatedDeltaNetAttention): return mixed_qkv_out, z_out, b_out, a_out + def _recoverssm_verify_qkv(self, mixed_qkv): + """Q/K/V for the RecoverSSM verify. + + The verify kernel reads token-strided heads, so for a 2-D [tokens, q|k|v] conv + output with a contiguous last dimension it gets views instead of copies. Any + other layout falls back to the contiguous copies of rearrange_mixed_qkv(). + Returns (query, key, value, is_view).""" + if mixed_qkv.ndim != 2 or mixed_qkv.stride(-1) != 1: + return (*self.rearrange_mixed_qkv(mixed_qkv), False) + q_dim = self.key_dim // self.tp_size + v_dim = self.value_dim // self.tp_size + num_tokens = mixed_qkv.shape[0] + query = mixed_qkv[:, :q_dim].view(1, num_tokens, -1, self.head_k_dim) + key = mixed_qkv[:, q_dim : 2 * q_dim].view(1, num_tokens, -1, self.head_k_dim) + value = mixed_qkv[:, 2 * q_dim : 2 * q_dim + v_dim].view( + 1, num_tokens, -1, self.head_v_dim + ) + return query, key, value, True + def rearrange_mixed_qkv(self, mixed_qkv): """Split packed qkv into contiguous (1, seq, heads, dim) tensors. @@ -1085,7 +1146,7 @@ class QwenGatedDeltaNetAttention(GatedDeltaNetAttention): dtype = qkv_or_qkvz.dtype num_k_heads = self.num_k_heads // self.tp_size num_v_heads = self.num_v_heads // self.tp_size - _, state_dtype = self.get_state_dtype() + state_dtype = self.get_state_dtype()[1] # 3 dtypes with the replay record # All kernels use BT = chunk_size, so a single pass with T = chunk_size # is sufficient to populate every autotuner cache. Mirror the real @@ -1343,7 +1404,11 @@ class QwenGatedDeltaNetAttention(GatedDeltaNetAttention): ], num_accepted_tokens=num_accepted_tokens, query_start_loc=spec_query_start_loc, - max_query_len=spec_state_indices_tensor.size(-1), + max_query_len=( + self.num_spec + 1 + if getattr(self, "use_gdn_recoverssm", False) + else spec_state_indices_tensor.size(-1) + ), validate_data=False, ) @@ -1380,7 +1445,17 @@ class QwenGatedDeltaNetAttention(GatedDeltaNetAttention): else: mixed_qkv_non_spec = None - query_spec, key_spec, value_spec = self.rearrange_mixed_qkv(mixed_qkv_spec) + if ( + spec_sequence_masks is not None + and mixed_qkv_spec is not None + and self.use_gdn_recoverssm + ): + query_spec, key_spec, value_spec, verify_views = ( + self._recoverssm_verify_qkv(mixed_qkv_spec) + ) + else: + verify_views = False + query_spec, key_spec, value_spec = self.rearrange_mixed_qkv(mixed_qkv_spec) # Split mixed non-spec-decode+prefill to process independently split_non_spec = ( @@ -1443,7 +1518,43 @@ class QwenGatedDeltaNetAttention(GatedDeltaNetAttention): # 2. Recurrent attention # 2.1: Process the multi-query part - if spec_sequence_masks is not None: + if spec_sequence_masks is not None and getattr( + self, "use_gdn_recoverssm", False + ): + # verify the window off the checkpoint; the commit after sampling + # writes the state once + from vllm.model_executor.layers.mamba.gdn.recoverssm_gdn import ( + gdn_recoverssm_verify, + ) + + _n = attn_metadata.num_spec_decodes + assert spec_query_start_loc is not None + assert spec_state_indices_tensor is not None + core_attn_out_spec = gdn_recoverssm_verify( + self.A_log, + a_spec, + b_spec, + self.dt_bias, + query_spec, + key_spec, + value_spec, + checkpoint_state=ssm_state, + replay_cache=self_kv_cache[2], + query_start_loc=spec_query_start_loc[: _n + 1], + state_indices=spec_state_indices_tensor[:_n, 0], + spec_query_len=self.num_spec + 1, + use_qk_l2norm_in_kernel=True, + # spec-only batch: write straight into the layer output + out=( + core_attn_out[:num_actual_tokens].unsqueeze(0) + if verify_views + and mixed_qkv_non_spec is None + and query_spec.shape[1] == num_actual_tokens + else None + ), + ) + last_recurrent_state = None + elif spec_sequence_masks is not None: core_attn_out_spec, last_recurrent_state = ( fused_sigmoid_gating_delta_rule_update( A_log=self.A_log, @@ -1558,7 +1669,8 @@ class QwenGatedDeltaNetAttention(GatedDeltaNetAttention): 0, non_spec_token_indx, core_attn_out_non_spec.squeeze(0) ) elif spec_sequence_masks is not None: - core_attn_out[:num_actual_tokens] = core_attn_out_spec.squeeze(0) + if core_attn_out_spec.data_ptr() != core_attn_out.data_ptr(): + core_attn_out[:num_actual_tokens] = core_attn_out_spec.squeeze(0) else: core_attn_out[:num_actual_tokens] = core_attn_out_non_spec.squeeze(0) diff --git a/vllm/model_executor/layers/mamba/gdn/recoverssm_gdn.py b/vllm/model_executor/layers/mamba/gdn/recoverssm_gdn.py new file mode 100644 index 0000000000..d347befcce --- /dev/null +++ b/vllm/model_executor/layers/mamba/gdn/recoverssm_gdn.py @@ -0,0 +1,695 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""RecoverSSM for Gated DeltaNet (Qwen-style GDN): speculative verify from one +checkpoint without per-draft state snapshots, and accepted-state recovery after +sampling. + +GDN port of the Kimi-K3 KDA RecoverSSM path +(vllm/models/kimi_k3/nvidia/ops/recoverssm.py). Differences from KDA: the gate is a +scalar per value head (g = -exp(A_log) * softplus(a + dt_bias)), q/k heads are grouped +over value heads (HV // H), and q/k are L2-normalised inside the kernel when requested. +The verify step caches, per token, the delta-rule correction c_t = beta_t * (v_t - S'_t +k_t) (V floats), the normalised key (K floats) and the log-decay g_t (1 float) in one +fp32 "replay" record. Recovery folds the accepted prefix: S_n = exp(sum_t g_t) S_0 + +sum_t exp(sum_{s>t} g_s) c_t k_t^T and writes the checkpoint once (plus the +block-boundary state in align mode).""" + +from collections.abc import Sequence +from dataclasses import dataclass +from typing import Any + +import torch + +from vllm.model_executor.layers.mamba.ops.recoverssm import ( + compact_conv_state_kernel, + prepare_commit_plan_kernel, +) +from vllm.model_executor.layers.mamba.recoverssm_utils import recoverssm_require +from vllm.triton_utils import tl, triton +from vllm.v1.attention.backends.utils import NULL_BLOCK_ID + + +def _require(cond: bool, msg: str) -> None: + recoverssm_require(cond, msg, "GDN RecoverSSM") + + +@triton.jit +def _gdn_recoverssm_verify_kernel( + q_ptr, + k_ptr, + v_ptr, + a_ptr, + b_ptr, + A_log_ptr, + dt_bias_ptr, + state_ptr, + replay_ptr, + out_ptr, + query_start_loc_ptr, + state_indices_ptr, + scale, + softplus_beta, + softplus_threshold, + null_block_id, + stride_q_token, + stride_k_token, + stride_v_token, + stride_a_token, + stride_b_token, + stride_state_block, + stride_state_head, + stride_state_v, + stride_state_k, + stride_replay_block, + stride_replay_head, + stride_replay_pos, + stride_replay_dim, + stride_out_token, + stride_qsl, + stride_si, + H: tl.constexpr, + HV: tl.constexpr, + K: tl.constexpr, + V: tl.constexpr, + BK: tl.constexpr, + BV: tl.constexpr, + SPEC_QUERY_LEN: tl.constexpr, + USE_QK_L2NORM: tl.constexpr, +): + pid_v = tl.program_id(0) + pid_b = tl.program_id(1) + pid_hv = tl.program_id(2) + pid_h = pid_hv // (HV // H) + bos = tl.load(query_start_loc_ptr + pid_b * stride_qsl).to(tl.int64) + eos = tl.load(query_start_loc_ptr + (pid_b + 1) * stride_qsl).to(tl.int64) + query_len = eos - bos + state_idx = tl.load(state_indices_ptr + pid_b * stride_si).to(tl.int64) + offs_k = tl.arange(0, BK) + offs_v = pid_v * BV + tl.arange(0, BV) + mask_k = offs_k < K + mask_v = offs_v < V + mask_state = mask_v[:, None] & mask_k[None, :] + if state_idx <= null_block_id: + for t in tl.static_range(SPEC_QUERY_LEN): + tl.store( + out_ptr + (bos + t) * stride_out_token + pid_hv * V + offs_v, + tl.zeros([BV], dtype=tl.float32), + mask=(t < query_len) & mask_v, + ) + return + state = tl.load( + state_ptr + + state_idx * stride_state_block + + pid_hv * stride_state_head + + offs_v[:, None] * stride_state_v + + offs_k[None, :] * stride_state_k, + mask=mask_state, + other=0.0, + ).to(tl.float32) + neg_A = -tl.exp(tl.load(A_log_ptr + pid_hv).to(tl.float32)) + dt_bias = tl.load(dt_bias_ptr + pid_hv).to(tl.float32) + rec_base = ( + replay_ptr + state_idx * stride_replay_block + pid_hv * stride_replay_head + ) + for t in tl.static_range(SPEC_QUERY_LEN): + valid = t < query_len + tok = bos + t + q = tl.load( + q_ptr + tok * stride_q_token + pid_h * K + offs_k, + mask=valid & mask_k, + other=0.0, + ).to(tl.float32) + k = tl.load( + k_ptr + tok * stride_k_token + pid_h * K + offs_k, + mask=valid & mask_k, + other=0.0, + ).to(tl.float32) + v = tl.load( + v_ptr + tok * stride_v_token + pid_hv * V + offs_v, + mask=valid & mask_v, + other=0.0, + ).to(tl.float32) + x = ( + tl.load(a_ptr + tok * stride_a_token + pid_hv, mask=valid, other=0.0).to( + tl.float32 + ) + + dt_bias + ) + sp = tl.where( + softplus_beta * x <= softplus_threshold, + (1 / softplus_beta) * tl.log(1 + tl.exp(softplus_beta * x)), + x, + ) + g = neg_A * sp + beta = tl.sigmoid( + tl.load(b_ptr + tok * stride_b_token + pid_hv, mask=valid, other=0.0).to( + tl.float32 + ) + ) + if USE_QK_L2NORM: + q = q * tl.rsqrt(tl.sum(q * q) + 1e-6) + k = k * tl.rsqrt(tl.sum(k * k) + 1e-6) + q = q * scale + new_state = state * tl.exp(g) + corr = (v - tl.sum(new_state * k[None, :], 1)) * beta + new_state = new_state + corr[:, None] * k[None, :] + state = tl.where(valid, new_state, state) + o = tl.sum(state * q[None, :], 1) + tl.store( + out_ptr + tok * stride_out_token + pid_hv * V + offs_v, + o.to(out_ptr.dtype.element_ty), + mask=valid & mask_v, + ) + rec = rec_base + t * stride_replay_pos + tl.store(rec + offs_v * stride_replay_dim, corr, mask=valid & mask_v) + if pid_v == 0: + tl.store(rec + (V + offs_k) * stride_replay_dim, k, mask=valid & mask_k) + tl.store(rec + (V + K) * stride_replay_dim, g, mask=valid) + + +def gdn_recoverssm_verify( + A_log: torch.Tensor, + a: torch.Tensor, + b: torch.Tensor, + dt_bias: torch.Tensor, + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + *, + checkpoint_state: torch.Tensor, + replay_cache: torch.Tensor, + query_start_loc: torch.Tensor, + state_indices: torch.Tensor, + spec_query_len: int, + scale: float | None = None, + softplus_beta: float = 1.0, + softplus_threshold: float = 20.0, + use_qk_l2norm_in_kernel: bool = True, + out: torch.Tensor | None = None, +) -> torch.Tensor: + """Verify a GDN speculative window against its checkpoint without modifying the + checkpoint. q, k: [1, T, H, K]; v: [1, T, HV, V]; a, b: [T, HV] (or [1, T, HV]); + checkpoint_state: [blocks, HV, V, K]; replay_cache: [blocks, HV, spec_query_len, V + + K + 1] fp32. Returns out [1, T, HV, V].""" + _require( + q.ndim == 4 and q.shape[0] == 1, "q must have shape [1, tokens, heads, dim]" + ) + _, total, H, K = q.shape + _require(k.shape == q.shape, "q and k shapes differ") + _require( + v.ndim == 4 and tuple(v.shape[:2]) == (1, total), + "v must have shape [1, tokens, value heads, dim]", + ) + HV, V = v.shape[-2], v.shape[-1] + _require( + H > 0 and HV >= H and HV % H == 0, + "value heads must be a positive multiple of the q/k heads", + ) + _require( + a.numel() == total * HV and b.numel() == total * HV, + "gate or beta shape is incompatible", + ) + a2 = a.reshape(total, HV) + b2 = b.reshape(total, HV) + _require( + a2.stride(1) == 1 and b2.stride(1) == 1, + "gate and beta heads must be contiguous", + ) + _require( + tuple(A_log.shape) == (HV,) and tuple(dt_bias.shape) == (HV,), + "A_log or dt_bias shape is incompatible", + ) + _require( + all(t_.stride(-1) == 1 and t_.stride(-2) == K for t_ in (q, k)), + "q and k heads must be contiguous", + ) + _require(v.stride(-1) == 1 and v.stride(-2) == V, "v heads must be contiguous") + _require(checkpoint_state.ndim == 4, "checkpoint must be four-dimensional") + nb = checkpoint_state.shape[0] + _require( + tuple(checkpoint_state.shape[1:]) == (HV, V, K), + "checkpoint shape is incompatible", + ) + _rec = (nb, HV, spec_query_len, V + K + 1) + _require(tuple(replay_cache.shape) == _rec, f"replay buffer needs shape {_rec}") + _require(replay_cache.dtype == torch.float32, "replay buffer must use float32") + _require(state_indices.ndim == 1, "state indices must be one-dimensional") + _require( + query_start_loc.ndim == 1 + and query_start_loc.shape[0] == state_indices.shape[0] + 1, + "query metadata is incompatible", + ) + _require( + total <= state_indices.shape[0] * spec_query_len, + "speculative decode input exceeds its activation capacity", + ) + _dev = q.device + _require( + all( + t_.device == _dev + for t_ in ( + k, + v, + a, + b, + A_log, + dt_bias, + checkpoint_state, + replay_cache, + query_start_loc, + state_indices, + ) + if t_ is not None + ) + and (out is None or out.device == _dev), + "inputs must be on the same device", + ) + if scale is None: + scale = K**-0.5 + if out is None: + out = torch.empty(1, total, HV, V, dtype=v.dtype, device=v.device) + _require( + tuple(out.shape) == (1, total, HV, V) and out.stride()[2:] == (V, 1), + "output shape or layout is incompatible", + ) + batch = state_indices.shape[0] + if total == 0 or batch == 0: + return out + BK = triton.next_power_of_2(K) + BV = min(triton.next_power_of_2(V), 32) + grid = (triton.cdiv(V, BV), batch, HV) + _gdn_recoverssm_verify_kernel[grid]( + q, + k, + v, + a2, + b2, + A_log, + dt_bias, + checkpoint_state, + replay_cache, + out, + query_start_loc, + state_indices, + scale, + softplus_beta, + softplus_threshold, + NULL_BLOCK_ID, + q.stride(1), + k.stride(1), + v.stride(1), + a2.stride(0), + b2.stride(0), + checkpoint_state.stride(0), + checkpoint_state.stride(1), + checkpoint_state.stride(2), + checkpoint_state.stride(3), + replay_cache.stride(0), + replay_cache.stride(1), + replay_cache.stride(2), + replay_cache.stride(3), + out.stride(1), + query_start_loc.stride(0), + state_indices.stride(0), + H=H, + HV=HV, + K=K, + V=V, + BK=BK, + BV=BV, + SPEC_QUERY_LEN=spec_query_len, + USE_QK_L2NORM=use_qk_l2norm_in_kernel, + num_warps=4, + num_stages=2, + ) + return out + + +@triton.jit +def _commit_gdn_state_kernel( + state_ref_ptr, + state_base_addrs_ptr, + state_block_strides_ptr, + replay_ref_ptr, + replay_base_addrs_ptr, + replay_block_strides_ptr, + state_indices_ptr, + commit_lens_ptr, + final_state_indices_ptr, + boundary_state_indices_ptr, + boundary_recovery_lens_ptr, + null_block_id, + stride_state_head, + stride_state_v, + stride_state_k, + stride_replay_head, + stride_replay_pos, + stride_replay_dim, + stride_si, + K: tl.constexpr, + V: tl.constexpr, + BK: tl.constexpr, + BV: tl.constexpr, + NUM_HEADS: tl.constexpr, + ALIGN_MODE: tl.constexpr, +): + pid_v = tl.program_id(0) + pid_b = tl.program_id(1) + pid_lh = tl.program_id(2) + pid_l = pid_lh // NUM_HEADS + pid_h = pid_lh % NUM_HEADS + src = tl.load(state_indices_ptr + pid_b * stride_si).to(tl.int64) + if src <= null_block_id: + return + n = tl.load(commit_lens_ptr + pid_b) + if n == 0: + return + final_idx = tl.load(final_state_indices_ptr + pid_b).to(tl.int64) + boundary_idx = tl.load(boundary_state_indices_ptr + pid_b).to(tl.int64) + boundary_len = tl.load(boundary_recovery_lens_ptr + pid_b) + if final_idx <= null_block_id: + return + state_ptr = tl.load(state_base_addrs_ptr + pid_l).to( + tl.pointer_type(state_ref_ptr.dtype.element_ty) + ) + sbs = tl.load(state_block_strides_ptr + pid_l) + replay_ptr = tl.load(replay_base_addrs_ptr + pid_l).to( + tl.pointer_type(replay_ref_ptr.dtype.element_ty) + ) + rbs = tl.load(replay_block_strides_ptr + pid_l) + rec_base = replay_ptr + src * rbs + pid_h * stride_replay_head + offs_k = tl.arange(0, BK) + offs_v = pid_v * BV + tl.arange(0, BV) + mask_k = offs_k < K + mask_v = offs_v < V + mask_state = mask_v[:, None] & mask_k[None, :] + sp = offs_v[:, None] * stride_state_v + offs_k[None, :] * stride_state_k + s0 = tl.load( + state_ptr + src * sbs + pid_h * stride_state_head + sp, + mask=mask_state, + other=0.0, + ).to(tl.float32) + fin_decay = 1.0 + fin_corr = tl.zeros([BV, BK], tl.float32) + bnd_decay = 1.0 + bnd_corr = tl.zeros([BV, BK], tl.float32) + for r in range(n): + t = n - r - 1 + rec = rec_base + t * stride_replay_pos + c = tl.load(rec + offs_v * stride_replay_dim, mask=mask_v, other=0.0) + kk = tl.load(rec + (V + offs_k) * stride_replay_dim, mask=mask_k, other=0.0) + g = tl.load(rec + (V + K) * stride_replay_dim) + upd = c[:, None] * kk[None, :] + dec = tl.exp(g) + fin_corr += upd * fin_decay + fin_decay *= dec + if ALIGN_MODE: + before = t < boundary_len + bnd_corr += tl.where(before, upd * bnd_decay, 0.0) + bnd_decay *= tl.where(before, dec, 1.0) + if ALIGN_MODE: + bs = s0 * bnd_decay + bnd_corr + tl.store( + state_ptr + boundary_idx * sbs + pid_h * stride_state_head + sp, + bs.to(state_ref_ptr.dtype.element_ty), + mask=mask_state & (boundary_idx > null_block_id), + ) + fs = s0 * fin_decay + fin_corr + tl.store( + state_ptr + final_idx * sbs + pid_h * stride_state_head + sp, + fs.to(state_ref_ptr.dtype.element_ty), + mask=mask_state, + ) + + +@dataclass +class GDNRecoverSSMCommitContext: + conv_states: tuple[torch.Tensor, ...] + conv_state_base_addrs: torch.Tensor + conv_state_block_strides: torch.Tensor + conv_state_dim_strides: torch.Tensor + conv_state_token_strides: torch.Tensor + conv_history_len: int + checkpoints: tuple[torch.Tensor, ...] + state_base_addrs: torch.Tensor + state_block_strides: torch.Tensor + replays: tuple[torch.Tensor, ...] + replay_base_addrs: torch.Tensor + replay_block_strides: torch.Tensor + commit_lens: torch.Tensor + final_state_indices: torch.Tensor + boundary_state_indices: torch.Tensor + boundary_recovery_lens: torch.Tensor + spec_query_len: int + + @classmethod + def from_tensors( + cls, + conv_states: Sequence[torch.Tensor], + checkpoints: Sequence[torch.Tensor], + replays: Sequence[torch.Tensor], + *, + spec_query_len: int, + max_num_reqs: int, + conv_dim_first: bool = True, + ) -> "GDNRecoverSSMCommitContext": + _require(len(checkpoints) > 0, "commit requires at least one layer") + _require( + len(conv_states) == len(checkpoints) == len(replays), + "conv, state and replay lists differ", + ) + if not conv_dim_first: + conv_states = [s.transpose(-1, -2) for s in conv_states] + ref = checkpoints[0] + _require(ref.ndim == 4, "checkpoint must be four-dimensional") + nb, HV, V, K = ref.shape + _dev = ref.device + for s in checkpoints: + _require( + s.shape == ref.shape + and s.dtype == ref.dtype + and s.stride()[1:] == ref.stride()[1:] + and s.device == _dev, + "layers need matching checkpoints", + ) + _rec = (nb, HV, spec_query_len, V + K + 1) + for r in replays: + _require( + tuple(r.shape) == _rec + and r.dtype == torch.float32 + and r.stride()[1:] == replays[0].stride()[1:] + and r.device == _dev, + f"layers need matching float32 replay buffers of shape {_rec}", + ) + conv_ref = conv_states[0] + _require( + conv_ref.ndim == 3 and conv_ref.shape[0] == nb, + "conv state must be [blocks, dim, window] with the checkpoint's " + "block count", + ) + for c in conv_states: + _require( + c.shape == conv_ref.shape + and c.dtype == conv_ref.dtype + and c.stride()[1:] == conv_ref.stride()[1:] + and c.device == _dev, + "layers need matching conv states", + ) + conv_dim, conv_len = conv_ref.shape[1:] + hist = conv_len - spec_query_len + 1 + _require(hist > 0, "conv state is shorter than its window") + dev = ref.device + addr = lambda ts: torch.tensor( + [t.data_ptr() for t in ts], dtype=torch.int64, device=dev + ) + bstr = lambda ts: torch.tensor( + [t.stride(0) for t in ts], dtype=torch.int64, device=dev + ) + z = lambda: torch.empty(max_num_reqs, dtype=torch.int32, device=dev) + return cls( + tuple(conv_states), + addr(conv_states), + bstr(conv_states), + torch.tensor( + [s.stride(1) for s in conv_states], dtype=torch.int64, device=dev + ), + torch.tensor( + [s.stride(2) for s in conv_states], dtype=torch.int64, device=dev + ), + hist, + tuple(checkpoints), + addr(checkpoints), + bstr(checkpoints), + tuple(replays), + addr(replays), + bstr(replays), + z(), + z(), + z(), + z(), + spec_query_len, + ) + + @classmethod + def create( + cls, layers: Sequence[Any], *, spec_query_len: int, max_num_reqs: int + ) -> "GDNRecoverSSMCommitContext": + from vllm.model_executor.layers.mamba.mamba_utils import is_conv_state_dim_first + + if any(len(layer.kv_cache) != 3 for layer in layers): + raise ValueError("GDN RecoverSSM pages must contain conv, state and replay") + return cls.from_tensors( + [layer.kv_cache[0] for layer in layers], + [layer.kv_cache[1] for layer in layers], + [layer.kv_cache[2] for layer in layers], + spec_query_len=spec_query_len, + max_num_reqs=max_num_reqs, + conv_dim_first=is_conv_state_dim_first(), + ) + + def commit( + self, + num_accepted_tokens: torch.Tensor, + state_indices: torch.Tensor, + query_start_loc: torch.Tensor, + request_indices: torch.Tensor | None = None, + block_table: torch.Tensor | None = None, + num_computed_tokens: torch.Tensor | None = None, + mamba_block_size: int | None = None, + commit_conv: bool = True, + ) -> None: + """Fold accepted GDN and convolution inputs into every layer's checkpoint.""" + batch = state_indices.shape[0] + if batch == 0: + return + _require(state_indices.ndim == 1, "state indices must be one-dimensional") + _require( + batch <= self.commit_lens.shape[0], "commit batch exceeds its plan capacity" + ) + _require( + query_start_loc.ndim == 1 and query_start_loc.shape[0] == batch + 1, + "commit metadata is incompatible", + ) + _require( + request_indices is None or request_indices.shape[0] >= batch, + "request mapping is too short", + ) + _require( + num_accepted_tokens.ndim == 1 + and (request_indices is not None or num_accepted_tokens.shape[0] >= batch), + "accepted-token counts are too short", + ) + _align = (block_table, num_computed_tokens, mamba_block_size) + _require( + all(x is None for x in _align) or all(x is not None for x in _align), + "align metadata is incomplete", + ) + _require( + mamba_block_size is None or mamba_block_size >= self.spec_query_len, + "align block size must cover one speculative window", + ) + _require( + block_table is None or block_table.ndim == 2, + "block table must be two-dimensional", + ) + _dev = self.checkpoints[0].device + _require( + all( + t_ is None or t_.device == _dev + for t_ in ( + num_accepted_tokens, + state_indices, + query_start_loc, + request_indices, + block_table, + num_computed_tokens, + ) + ), + "commit inputs must be on the same device", + ) + bt_stride = (0, 0) if block_table is None else block_table.stride() + prepare_commit_plan_kernel[(batch,)]( + num_accepted_tokens, + request_indices, + state_indices, + query_start_loc, + block_table, + num_computed_tokens, + self.commit_lens, + self.final_state_indices, + self.boundary_state_indices, + self.boundary_recovery_lens, + NULL_BLOCK_ID, + mamba_block_size or 1, + block_table.shape[1] if block_table is not None else 1, + num_accepted_tokens.stride(0), + request_indices.stride(0) if request_indices is not None else 0, + state_indices.stride(0), + query_start_loc.stride(0), + bt_stride[0], + bt_stride[1], + 0 if num_computed_tokens is None else num_computed_tokens.stride(0), + SPEC_QUERY_LEN=self.spec_query_len, + num_warps=1, + ) + num_layers = len(self.checkpoints) + if commit_conv: + conv_ref = self.conv_states[0] + conv_dim = conv_ref.shape[1] + compact_conv_state_kernel[(triton.cdiv(conv_dim, 256), batch, num_layers)]( + conv_ref, + self.conv_state_base_addrs, + self.conv_state_block_strides, + self.conv_state_dim_strides, + self.conv_state_token_strides, + state_indices, + self.commit_lens, + self.final_state_indices, + self.boundary_state_indices, + self.boundary_recovery_lens, + NULL_BLOCK_ID, + conv_dim, + self.conv_history_len, + state_indices.stride(0), + BLOCK_D=256, + BLOCK_HISTORY=triton.next_power_of_2(self.conv_history_len), + ALIGN_MODE=block_table is not None, + num_warps=4, + ) + ref = self.checkpoints[0] + _, HV, V, K = ref.shape + BK = triton.next_power_of_2(K) + BV = min(triton.next_power_of_2(V), 32) + rr = self.replays[0] + _commit_gdn_state_kernel[(triton.cdiv(V, BV), batch, num_layers * HV)]( + ref, + self.state_base_addrs, + self.state_block_strides, + rr, + self.replay_base_addrs, + self.replay_block_strides, + state_indices, + self.commit_lens, + self.final_state_indices, + self.boundary_state_indices, + self.boundary_recovery_lens, + NULL_BLOCK_ID, + ref.stride(1), + ref.stride(2), + ref.stride(3), + rr.stride(1), + rr.stride(2), + rr.stride(3), + state_indices.stride(0), + K=K, + V=V, + BK=BK, + BV=BV, + NUM_HEADS=HV, + ALIGN_MODE=block_table is not None, + num_warps=4, + num_stages=2, + ) + + +__all__ = ["GDNRecoverSSMCommitContext", "gdn_recoverssm_verify"] diff --git a/vllm/model_executor/models/interfaces.py b/vllm/model_executor/models/interfaces.py index 7901a136f5..6653aafbe9 100644 --- a/vllm/model_executor/models/interfaces.py +++ b/vllm/model_executor/models/interfaces.py @@ -85,6 +85,7 @@ MambaStateShapes: TypeAlias = ( | tuple[tuple[int, int, int]] | tuple[tuple[int, int], tuple[int, int]] | tuple[tuple[int, int], tuple[int, int, int]] + | tuple[tuple[int, int], tuple[int, int, int], tuple[int, int, int]] | tuple[ tuple[int, int], tuple[int, int, int], diff --git a/vllm/models/qwen4_exp/nvidia/model.py b/vllm/models/qwen4_exp/nvidia/model.py index 2e38339655..229a5a4a53 100644 --- a/vllm/models/qwen4_exp/nvidia/model.py +++ b/vllm/models/qwen4_exp/nvidia/model.py @@ -32,11 +32,13 @@ from vllm.model_executor.layers.vocab_parallel_embedding import ( from vllm.model_executor.models.interfaces import ( HasInnerState, IsHybrid, + MambaStateShapes, MixtureOfExperts, MultiModalEmbeddings, SupportsLoRA, SupportsMRoPE, SupportsPP, + SupportsReplaySSM, _require_is_multimodal, ) from vllm.model_executor.models.qwen3_5 import ( @@ -641,6 +643,7 @@ class Qwen4ExpForCausalLM( SupportsPP, Qwen4ExpMixtureOfExperts, IsHybrid, + SupportsReplaySSM, ): packed_modules_mapping = { "qkv_proj": ["q_proj", "k_proj", "v_proj"], @@ -746,20 +749,38 @@ class Qwen4ExpForCausalLM( num_spec=num_spec, ) + @classmethod + def _uses_gdn_recoverssm(cls, vllm_config: VllmConfig) -> bool: + """One predicate for the replay record's shape AND dtype (they must + agree for the MambaSpec).""" + num_spec = ( + vllm_config.speculative_config.num_speculative_tokens + if vllm_config.speculative_config + else 0 + ) + from vllm.model_executor.layers.mamba.recoverssm_utils import uses_recoverssm + + return uses_recoverssm(vllm_config.cache_config, num_spec or 0) + @classmethod def get_gdn_mamba_state_dtype_from_config( cls, vllm_config: VllmConfig - ) -> tuple[torch.dtype, torch.dtype]: - return MambaStateDtypeCalculator.gated_delta_net_state_dtype( - vllm_config.model_config.dtype, - vllm_config.cache_config.mamba_cache_dtype, - vllm_config.cache_config.mamba_ssm_cache_dtype, + ) -> tuple[torch.dtype, ...]: + dtypes: tuple[torch.dtype, ...] = ( + MambaStateDtypeCalculator.gated_delta_net_state_dtype( + vllm_config.model_config.dtype, + vllm_config.cache_config.mamba_cache_dtype, + vllm_config.cache_config.mamba_ssm_cache_dtype, + ) ) + if cls._uses_gdn_recoverssm(vllm_config): # replay record, fp32 + dtypes = (*dtypes, torch.float32) + return dtypes @classmethod def get_gdn_mamba_state_shape_from_config( cls, vllm_config: VllmConfig - ) -> tuple[tuple[int, int], tuple[int, int]]: + ) -> MambaStateShapes: parallel_config = vllm_config.parallel_config hf_config = vllm_config.model_config.hf_text_config tp_size = parallel_config.tensor_parallel_size @@ -768,7 +789,7 @@ class Qwen4ExpForCausalLM( if vllm_config.speculative_config else 0 ) - return MambaStateShapeCalculator.gated_delta_net_state_shape( + shapes = MambaStateShapeCalculator.gated_delta_net_state_shape( tp_size, hf_config.linear_num_key_heads, hf_config.linear_num_value_heads, @@ -777,18 +798,30 @@ class Qwen4ExpForCausalLM( hf_config.linear_conv_kernel_dim, num_spec, ) + if cls._uses_gdn_recoverssm( + vllm_config + ): # replay record [HV, spec_query_len, V + K + 1] + shapes = ( + *shapes, + ( + hf_config.linear_num_value_heads // tp_size, + num_spec + 1, + hf_config.linear_value_head_dim + hf_config.linear_key_head_dim + 1, + ), + ) + return shapes @classmethod def get_mamba_state_dtype_from_config( cls, vllm_config: VllmConfig, - ) -> tuple[torch.dtype, torch.dtype]: + ) -> tuple[torch.dtype, ...]: return cls.get_gdn_mamba_state_dtype_from_config(vllm_config) @classmethod def get_mamba_state_shape_from_config( cls, vllm_config: VllmConfig - ) -> tuple[tuple[int, int], tuple[int, int]]: + ) -> MambaStateShapes: return cls.get_gdn_mamba_state_shape_from_config(vllm_config) @classmethod @@ -876,6 +909,7 @@ class Qwen4ExpForConditionalGeneration( Qwen3_5ForConditionalGeneration, HasInnerState, Qwen4ExpMixtureOfExperts, + SupportsReplaySSM, ): """Qwen3-VL vision tower backed by the Qwen4Exp language model.""" @@ -1060,14 +1094,15 @@ class Qwen4ExpForConditionalGeneration( def get_mamba_state_dtype_from_config( cls, vllm_config: VllmConfig, - ) -> tuple[torch.dtype, torch.dtype]: + ) -> tuple[torch.dtype, ...]: return Qwen4ExpForCausalLM.get_mamba_state_dtype_from_config(vllm_config) @classmethod - def get_mamba_state_shape_from_config( + def get_mamba_state_shape_from_config( # type: ignore[override] cls, vllm_config: VllmConfig, - ) -> tuple[tuple[int, int], tuple[int, int]]: + ) -> MambaStateShapes: + # GDN RecoverSSM adds a third state (the replay record) to Qwen3.5's two. return Qwen4ExpForCausalLM.get_mamba_state_shape_from_config(vllm_config) @classmethod diff --git a/vllm/models/qwen4_exp/nvidia/mtp.py b/vllm/models/qwen4_exp/nvidia/mtp.py index af3bcd4659..6399b28f2a 100644 --- a/vllm/models/qwen4_exp/nvidia/mtp.py +++ b/vllm/models/qwen4_exp/nvidia/mtp.py @@ -32,7 +32,7 @@ from vllm.model_executor.layers.vocab_parallel_embedding import ( VocabParallelEmbedding, ) from vllm.model_executor.model_loader.utils import configure_quant_config -from vllm.model_executor.models.interfaces import SupportsPP +from vllm.model_executor.models.interfaces import SupportsPP, SupportsReplaySSM from vllm.model_executor.models.qwen3_5 import Qwen3_5Model from vllm.model_executor.models.utils import ( AutoWeightsLoader, @@ -359,7 +359,7 @@ class Qwen4ExpMultiTokenPredictor(nn.Module): return loader.load_weights(weights, mapper=mapper) -class Qwen4ExpMTP(nn.Module, SupportsPP, Qwen4ExpMixtureOfExperts): +class Qwen4ExpMTP(nn.Module, SupportsPP, Qwen4ExpMixtureOfExperts, SupportsReplaySSM): packed_modules_mapping = { "qkv_proj": ["q_proj", "k_proj", "v_proj"], "gate_up_proj": ["gate_proj", "up_proj"], diff --git a/vllm/models/qwen4_exp/nvidia/ple_layer.py b/vllm/models/qwen4_exp/nvidia/ple_layer.py index b04cdaad21..c196c68b10 100644 --- a/vllm/models/qwen4_exp/nvidia/ple_layer.py +++ b/vllm/models/qwen4_exp/nvidia/ple_layer.py @@ -175,6 +175,14 @@ class Qwen4ExpPLELayer(nn.Module, MambaBase): return True def get_attn_backend(self) -> type[PleShortConvAttentionBackend]: + from vllm.model_executor.layers.mamba.recoverssm_utils import uses_recoverssm + + if uses_recoverssm(self.cache_config, self.num_spec_tokens): + from vllm.v1.attention.backends.ple_recoverssm import ( + PleRecoverSSMAttentionBackend, + ) + + return PleRecoverSSMAttentionBackend return PleShortConvAttentionBackend def get_state_dtype(self) -> tuple[torch.dtype, ...]: diff --git a/vllm/v1/attention/backends/gdn_recoverssm.py b/vllm/v1/attention/backends/gdn_recoverssm.py new file mode 100644 index 0000000000..a45142ab3c --- /dev/null +++ b/vllm/v1/attention/backends/gdn_recoverssm.py @@ -0,0 +1,181 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""GDN attention metadata for RecoverSSM speculative decode (Qwen-style Gated DeltaNet). + +Wraps the generic GDN builder: every active decode row runs through the spec path (a +draft-less step is a T=1 window), the spec state indices keep only the checkpoint +column, num_accepted_tokens is always 1 (the checkpoint and the compacted conv window +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 typing import Any + +import torch + +from vllm.logger import init_logger +from vllm.v1.attention.backend import AttentionCGSupport, CommonAttentionMetadata +from vllm.v1.attention.backends.gdn_attn import ( + GDNAttentionBackend, + GDNAttentionMetadata, + GDNAttentionMetadataBuilder, +) +from vllm.v1.attention.backends.recoverssm_metadata import ( + RecoverSSMMetadata, + RecoverSSMPostprocessMetadata, + recoverssm_request_indices, + recoverssm_spec_rows, +) + +logger = init_logger(__name__) + + +@dataclass +class GDNRecoverSSMCommitMetadata: + state_indices: torch.Tensor # [num_spec_decodes, 1] + query_start_loc: torch.Tensor # [num_spec_decodes + 1] + request_indices: ( + torch.Tensor | None + ) # batch rows of the spec requests (None: rows 0..n-1) + block_table: torch.Tensor | None # align mode + num_computed_tokens: torch.Tensor | None + block_size: int | None + + +@dataclass +class GDNRecoverSSMMetadata(GDNAttentionMetadata, RecoverSSMMetadata): + recoverssm_commit: GDNRecoverSSMCommitMetadata | None = None + recoverssm_context: Any = field(default=None, repr=False, compare=False) + + def commit_recoverssm_state( + self, num_accepted_tokens: torch.Tensor + ) -> RecoverSSMPostprocessMetadata | None: + c = self.recoverssm_commit + if c is None: + return None + n = self.num_spec_decodes + self.recoverssm_context.commit( + num_accepted_tokens, + c.state_indices[:n, 0], + c.query_start_loc[: n + 1], + request_indices=c.request_indices, + block_table=c.block_table, + num_computed_tokens=c.num_computed_tokens, + mamba_block_size=c.block_size, + ) + if c.block_table is None: + return None + assert c.block_size is not None and c.num_computed_tokens is not None + return RecoverSSMPostprocessMetadata( + num_spec_decodes=n, + request_indices=c.request_indices, + block_table=c.block_table, + num_computed_tokens=c.num_computed_tokens, + block_size=c.block_size, + ) + + +class GDNRecoverSSMMetadataBuilder(GDNAttentionMetadataBuilder): + # The verify/commit metadata is not staged for FULL graph replay yet; declaring + # NEVER makes the runner fall back to PIECEWISE (attention compiled piecewise). + _cudagraph_support = AttentionCGSupport.NEVER + + def __init__(self, kv_cache_spec, layer_names, vllm_config, device) -> None: + super().__init__(kv_cache_spec, layer_names, vllm_config, device) + # The builder may be created before the runner resolves the CUDA graph mode. + self.use_full_cuda_graph = False + max_reqs = vllm_config.scheduler_config.max_num_seqs + self._ones = torch.ones(max_reqs, dtype=torch.int32, device=device) + self._context: Any = None + self._logged = False + + def _get_context(self): + if self._context is None: + from vllm.model_executor.layers.mamba.gdn.recoverssm_gdn import ( + GDNRecoverSSMCommitContext, + ) + + fc = self.vllm_config.compilation_config.static_forward_context + layers = [fc[name] for name in self.layer_names] + self._context = GDNRecoverSSMCommitContext.create( + layers, + spec_query_len=1 + self.num_spec, + max_num_reqs=self.vllm_config.scheduler_config.max_num_seqs, + ) + logger.info( + "GDN RecoverSSM commit context: %d layers, spec_query_len %d", + len(layers), + 1 + self.num_spec, + ) + return self._context + + def build( # type: ignore[override] + self, + common_prefix_len: int, + common_attn_metadata: CommonAttentionMetadata, + num_accepted_tokens: torch.Tensor | None = None, + num_decode_draft_tokens_cpu: torch.Tensor | None = None, + fast_build: bool = False, + ) -> "GDNRecoverSSMMetadata": + m = common_attn_metadata + rows = None + if self.use_spec_decode and num_decode_draft_tokens_cpu is not None: + if m.is_prefilling is None or m.is_prefilling.device.type != "cpu": + raise ValueError("GDN RecoverSSM needs the CPU is_prefilling mask") + # Positive counts only defeat the base builder's "no drafts -> regular + # decode" shortcut: RecoverSSM keeps every post-prefill row on the spec path + # (extended conv window). + num_decode_draft_tokens_cpu, rows = recoverssm_spec_rows( + m.is_prefilling, + m.query_start_loc_cpu, + num_decode_draft_tokens_cpu, + m.num_reqs, + ) + meta = super().build( + common_prefix_len, + m, + num_accepted_tokens=num_accepted_tokens, + num_decode_draft_tokens_cpu=num_decode_draft_tokens_cpu, + fast_build=fast_build, + ) + base = {f.name: getattr(meta, f.name) for f in fields(meta)} + commit = None + if meta.num_spec_decodes > 0: + n = meta.num_spec_decodes + assert meta.spec_state_indices_tensor is not None + base["spec_state_indices_tensor"] = meta.spec_state_indices_tensor[:, :1] + base["num_accepted_tokens"] = self._ones[:n] + if rows is None: + raise ValueError("GDN RecoverSSM: spec rows without draft counts") + request_indices = recoverssm_request_indices( + rows, n, m.query_start_loc.device + ) + align = self.vllm_config.cache_config.mamba_cache_mode == "align" + commit = GDNRecoverSSMCommitMetadata( + state_indices=base["spec_state_indices_tensor"], + query_start_loc=meta.spec_query_start_loc, + request_indices=request_indices, + block_table=m.block_table_tensor if align else None, + num_computed_tokens=m.compute_num_computed_tokens() if align else None, + block_size=self.kv_cache_spec.block_size if align else None, + ) + if not self._logged: + self._logged = True + logger.info( + "GDN RecoverSSM path taken: %d spec rows, align=%s", n, align + ) + return GDNRecoverSSMMetadata( + **base, + recoverssm_commit=commit, + recoverssm_context=self._get_context() if commit is not None else None, + ) + + +class GDNRecoverSSMAttentionBackend(GDNAttentionBackend): + @staticmethod + def get_name() -> str: + return "GDN_RECOVERSSM" + + @staticmethod + def get_builder_cls() -> type[GDNRecoverSSMMetadataBuilder]: + return GDNRecoverSSMMetadataBuilder diff --git a/vllm/v1/attention/backends/mamba_attn.py b/vllm/v1/attention/backends/mamba_attn.py index 63f98f1786..4230dc39da 100644 --- a/vllm/v1/attention/backends/mamba_attn.py +++ b/vllm/v1/attention/backends/mamba_attn.py @@ -112,7 +112,12 @@ class BaseMambaAttentionMetadataBuilder(AttentionMetadataBuilder[M], abc.ABC): self.compilation_config = vllm_config.compilation_config self.num_spec_tokens: int = vllm_config.num_speculative_tokens self.use_spec_decode = self.num_spec_tokens > 0 - self.use_replayssm = vllm_config.cache_config.use_replayssm + # RecoverSSM also rides on --use-replayssm but keeps its own metadata; the + # Mamba2 ReplaySSM ring-buffer workspaces below are not used then. + self.use_replayssm = ( + vllm_config.cache_config.use_replayssm + and not vllm_config.cache_config.use_recoverssm + ) self.replayssm_buffer_len = vllm_config.cache_config.replayssm_buffer_len self.use_flashinfer_replayssm = ( self.use_replayssm diff --git a/vllm/v1/attention/backends/ple_recoverssm.py b/vllm/v1/attention/backends/ple_recoverssm.py new file mode 100644 index 0000000000..7402001fc0 --- /dev/null +++ b/vllm/v1/attention/backends/ple_recoverssm.py @@ -0,0 +1,326 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""RecoverSSM protocol for the Qwen4Exp PLE short-conv layer. + +When the GDN layers run RecoverSSM, the runner resets the shared num_accepted_tokens to +1 after each commit (align mode). The PLE dilated short conv keeps an extended window in +one state block and reads it at num_accepted - 1, so it must follow the same protocol: +every active decode row takes the spec path, the conv reads at offset 0, and after +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 typing import Any + +import torch + +from vllm.logger import init_logger +from vllm.model_executor.layers.mamba.ops.recoverssm import ( + compact_conv_state_kernel, + prepare_commit_plan_kernel, +) +from vllm.model_executor.layers.mamba.recoverssm_utils import recoverssm_require +from vllm.triton_utils import triton +from vllm.v1.attention.backend import AttentionCGSupport +from vllm.v1.attention.backends.recoverssm_metadata import ( + RecoverSSMMetadata, + RecoverSSMPostprocessMetadata, + recoverssm_request_indices, + recoverssm_spec_rows, +) +from vllm.v1.attention.backends.short_conv_attn import ( + PleShortConvAttentionBackend, + PleShortConvAttentionMetadata, + PleShortConvAttentionMetadataBuilder, +) +from vllm.v1.attention.backends.utils import NULL_BLOCK_ID + +logger = init_logger(__name__) + + +def _require(cond: bool, msg: str) -> None: + recoverssm_require(cond, msg, "PLE RecoverSSM") + + +class _PleConvCommit: + def __init__(self, conv_states, spec_query_len: int, max_num_reqs: int): + from vllm.model_executor.layers.mamba.mamba_utils import is_conv_state_dim_first + + _require(len(conv_states) > 0, "PLE commit requires at least one layer") + if not is_conv_state_dim_first(): + conv_states = [ + s.transpose(-1, -2) for s in conv_states + ] # -> [blocks, C, W] + self.conv_states = list(conv_states) + ref = self.conv_states[0] + _require(ref.ndim == 3, "PLE conv state must be [blocks, dim, window]") + dev = ref.device + for s in self.conv_states: + _require( + s.shape == ref.shape + and s.dtype == ref.dtype + and s.stride()[1:] == ref.stride()[1:] + and s.device == dev, + "PLE layers need matching conv states", + ) + self.conv_dim, conv_len = ref.shape[1], ref.shape[2] + self.hist = conv_len - spec_query_len + 1 + _require(self.hist > 0, "PLE conv state is shorter than its window") + self.spec_query_len = spec_query_len + t = lambda xs: torch.tensor(xs, dtype=torch.int64, device=dev) + self.base = t([s.data_ptr() for s in self.conv_states]) + self.bstride = t([s.stride(0) for s in self.conv_states]) + self.dstride = t([s.stride(1) for s in self.conv_states]) + self.tstride = t([s.stride(2) for s in self.conv_states]) + z = lambda: torch.empty(max_num_reqs, dtype=torch.int32, device=dev) + self.commit_lens, self.final_idx, self.bnd_idx, self.bnd_len = ( + z(), + z(), + z(), + z(), + ) + + def commit( + self, + num_accepted_tokens, + state_indices, + query_start_loc, + request_indices=None, + block_table=None, + num_computed_tokens=None, + mamba_block_size=None, + ): + batch = state_indices.shape[0] + if batch == 0: + return + _require(state_indices.ndim == 1, "state indices must be one-dimensional") + _require( + batch <= self.commit_lens.shape[0], + "PLE commit batch exceeds its plan capacity", + ) + _require( + query_start_loc.ndim == 1 and query_start_loc.shape[0] == batch + 1, + "PLE commit metadata is incompatible", + ) + _require( + request_indices is None or request_indices.shape[0] >= batch, + "PLE request mapping is too short", + ) + _require( + num_accepted_tokens.ndim == 1 + and (request_indices is not None or num_accepted_tokens.shape[0] >= batch), + "PLE accepted-token counts are too short", + ) + _align = (block_table, num_computed_tokens, mamba_block_size) + _require( + all(x is None for x in _align) or all(x is not None for x in _align), + "PLE align metadata is incomplete", + ) + _require( + mamba_block_size is None or mamba_block_size >= self.spec_query_len, + "PLE align block size must cover one speculative window", + ) + _require( + block_table is None or block_table.ndim == 2, + "PLE block table must be two-dimensional", + ) + _dev = self.conv_states[0].device + _require( + all( + t_ is None or t_.device == _dev + for t_ in ( + num_accepted_tokens, + state_indices, + query_start_loc, + request_indices, + block_table, + num_computed_tokens, + ) + ), + "PLE commit inputs must be on the same device", + ) + bt = (0, 0) if block_table is None else block_table.stride() + prepare_commit_plan_kernel[(batch,)]( + num_accepted_tokens, + request_indices, + state_indices, + query_start_loc, + block_table, + num_computed_tokens, + self.commit_lens, + self.final_idx, + self.bnd_idx, + self.bnd_len, + NULL_BLOCK_ID, + mamba_block_size or 1, + block_table.shape[1] if block_table is not None else 1, + num_accepted_tokens.stride(0), + request_indices.stride(0) if request_indices is not None else 0, + state_indices.stride(0), + query_start_loc.stride(0), + bt[0], + bt[1], + 0 if num_computed_tokens is None else num_computed_tokens.stride(0), + SPEC_QUERY_LEN=self.spec_query_len, + num_warps=1, + ) + compact_conv_state_kernel[ + (triton.cdiv(self.conv_dim, 256), batch, len(self.conv_states)) + ]( + self.conv_states[0], + self.base, + self.bstride, + self.dstride, + self.tstride, + state_indices, + self.commit_lens, + self.final_idx, + self.bnd_idx, + self.bnd_len, + NULL_BLOCK_ID, + self.conv_dim, + self.hist, + state_indices.stride(0), + BLOCK_D=256, + BLOCK_HISTORY=triton.next_power_of_2(self.hist), + ALIGN_MODE=block_table is not None, + num_warps=4, + ) + + +@dataclass +class PleRecoverSSMCommitMetadata: + state_indices: torch.Tensor # [n] + query_start_loc: torch.Tensor # [n + 1] + request_indices: torch.Tensor | None + block_table: torch.Tensor | None + num_computed_tokens: torch.Tensor | None + block_size: int | None + + +@dataclass +class PleRecoverSSMMetadata(PleShortConvAttentionMetadata, RecoverSSMMetadata): + recoverssm_commit: PleRecoverSSMCommitMetadata | None = None + recoverssm_context: Any = field(default=None, repr=False, compare=False) + + def commit_recoverssm_state(self, num_accepted_tokens): + c = self.recoverssm_commit + if c is None: + return None + n = self.num_spec_decodes + self.recoverssm_context.commit( + num_accepted_tokens, + c.state_indices[:n], + c.query_start_loc[: n + 1], + request_indices=c.request_indices, + block_table=c.block_table, + num_computed_tokens=c.num_computed_tokens, + mamba_block_size=c.block_size, + ) + if c.block_table is None: + return None + assert c.block_size is not None and c.num_computed_tokens is not None + return RecoverSSMPostprocessMetadata( + num_spec_decodes=n, + request_indices=c.request_indices, + block_table=c.block_table, + num_computed_tokens=c.num_computed_tokens, + block_size=c.block_size, + ) + + +class PleRecoverSSMMetadataBuilder(PleShortConvAttentionMetadataBuilder): + # Same as the GDN RecoverSSM builder: PIECEWISE only for now. + _cudagraph_support = AttentionCGSupport.NEVER + + def __init__(self, kv_cache_spec, layer_names, vllm_config, device) -> None: + super().__init__(kv_cache_spec, layer_names, vllm_config, device) + self.use_full_cuda_graph = False + self._ones = torch.ones( + vllm_config.scheduler_config.max_num_seqs, dtype=torch.int32, device=device + ) + self._ctx: _PleConvCommit | None = None + self._logged = False + + def _get_ctx(self): + 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] + ctx = _PleConvCommit( + convs, 1 + self.num_spec, self.vllm_config.scheduler_config.max_num_seqs + ) + self._ctx = ctx + logger.info( + "PLE RecoverSSM commit context: %d layers, history %d", + len(convs), + ctx.hist, + ) + return self._ctx + + def build( + self, + common_prefix_len, + common_attn_metadata, + fast_build=False, + *, + num_accepted_tokens=None, + num_decode_draft_tokens_cpu=None, + **kwargs, + ): # type: ignore[override] + m = common_attn_metadata + rows = None + if self.use_spec_decode and num_decode_draft_tokens_cpu is not None: + if m.is_prefilling is None or m.is_prefilling.device.type != "cpu": + raise ValueError("PLE RecoverSSM needs the CPU is_prefilling mask") + # the same classification as the GDN builder, so both commits map the same + # requests + num_decode_draft_tokens_cpu, rows = recoverssm_spec_rows( + m.is_prefilling, + m.query_start_loc_cpu, + num_decode_draft_tokens_cpu, + m.num_reqs, + ) + meta = super().build( + common_prefix_len, + m, + fast_build, + num_accepted_tokens=num_accepted_tokens, + num_decode_draft_tokens_cpu=num_decode_draft_tokens_cpu, + **kwargs, + ) + if ( + not isinstance(meta, PleShortConvAttentionMetadata) + or meta.num_spec_decodes == 0 + ): + return meta + base = {f.name: getattr(meta, f.name) for f in fields(meta)} + n = meta.num_spec_decodes + base["num_accepted_tokens"] = self._ones[:n] + if rows is None: + raise ValueError("PLE RecoverSSM: spec rows without draft counts") + req_idx = recoverssm_request_indices(rows, n, m.query_start_loc.device) + align = self.vllm_config.cache_config.mamba_cache_mode == "align" + commit = PleRecoverSSMCommitMetadata( + state_indices=meta.spec_state_indices_tensor, + query_start_loc=meta.spec_query_start_loc, + request_indices=req_idx, + block_table=m.block_table_tensor if align else None, + num_computed_tokens=m.compute_num_computed_tokens() if align else None, + block_size=self.kv_cache_spec.block_size if align else None, + ) + if not self._logged: + self._logged = True + logger.info("PLE RecoverSSM path taken: %d spec rows, align=%s", n, align) + return PleRecoverSSMMetadata( + **base, recoverssm_commit=commit, recoverssm_context=self._get_ctx() + ) + + +class PleRecoverSSMAttentionBackend(PleShortConvAttentionBackend): + @staticmethod + def get_name() -> str: + return "PLE_SHORT_CONV_RECOVERSSM" + + @staticmethod + def get_builder_cls() -> type[PleRecoverSSMMetadataBuilder]: + return PleRecoverSSMMetadataBuilder From 5567cc1b25c6fe79b99d228068017e371efc2570 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?J=C3=BCrgen=20Schmied?= Date: Sun, 27 Sep 2026 11:06:54 +0200 Subject: [PATCH 3/3] [Mamba] RecoverSSM align: keep a boundary-ending state in its block MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit A speculative window whose accepted tokens end exactly on a recurrent block boundary selected the following block-table column for the committed state (commit plan) and recorded it as the running column (postprocess). That column can still be unallocated, so the next step reads an uninitialized page. Select the block that holds the last committed token, matching the (n - 1) // block convention used elsewhere in align mode. Applies to Kimi-K3 KDA and to the Qwen GDN/PLE RecoverSSM paths that share the plan. Fix and test from local-inference-lab/vllm a02f585d (Martin Vit, root cause reported by astra); the GDN/PLE align tests now expect the same convention. Signed-off-by: Jürgen Schmied Co-authored-by: Martin Vit Co-authored-by: Claude --- tests/kernels/mamba/test_recoverssm_gdn.py | 9 +- tests/kernels/mamba/test_recoverssm_ple.py | 2 +- .../kimi_k3/test_recoverssm_boundary.py | 129 ++++++++++++++++++ .../layers/mamba/ops/recoverssm.py | 6 +- vllm/v1/worker/gpu/model_states/recoverssm.py | 4 +- 5 files changed, 143 insertions(+), 7 deletions(-) create mode 100644 tests/models/kimi_k3/test_recoverssm_boundary.py diff --git a/tests/kernels/mamba/test_recoverssm_gdn.py b/tests/kernels/mamba/test_recoverssm_gdn.py index 2488f2bd5c..08dcb4f8ac 100644 --- a/tests/kernels/mamba/test_recoverssm_gdn.py +++ b/tests/kernels/mamba/test_recoverssm_gdn.py @@ -299,14 +299,15 @@ def test_commit_maps_noncontiguous_requests_in_align_mode(): for i, r in enumerate(req_rows.tolist()): n = min(int(acc[r]), qlens[i]) c0 = int(nc[r]) - final_blk = int(bt[r, (c0 + n) // bs]) + final_blk = int(bt[r, max(c0 + n - 1, 0) // bs]) tok = int(inp["query_start_loc"][i]) + n - 1 torch.testing.assert_close(ckpt[final_blk], states[tok], rtol=1e-5, atol=1e-6) # (num_computed, accepted, block size): boundary after 1 / after 2 / ending exactly on # it / crossing and continuing / accepted = 1 without crossing / the full window -# crossing after 3 +# crossing after 3. A window ending exactly on a boundary keeps its state in the block +# holding its last token: the next block may still be unallocated. ALIGN_CASES = [ (3, 2, 4), (2, 3, 4), @@ -338,7 +339,7 @@ def test_align_boundary_and_final_state_match_native(nc, accepted, bs): mamba_block_size=bs, ) next_boundary = (nc // bs + 1) * bs - final_blk = int(bt[0, (nc + accepted) // bs]) + final_blk = int(bt[0, max(nc + accepted - 1, 0) // bs]) if ( nc + accepted >= next_boundary ): # crosses: the boundary state goes into the source block @@ -585,7 +586,7 @@ def test_align_boundary_bf16_state(nc, accepted, bs): mamba_block_size=bs, ) next_boundary = (nc // bs + 1) * bs - final_blk = int(bt[0, (nc + accepted) // bs]) + final_blk = int(bt[0, max(nc + accepted - 1, 0) // bs]) if nc + accepted >= next_boundary: torch.testing.assert_close( ckpt[int(src)].float(), diff --git a/tests/kernels/mamba/test_recoverssm_ple.py b/tests/kernels/mamba/test_recoverssm_ple.py index ec11ad3405..6e2e1f0526 100644 --- a/tests/kernels/mamba/test_recoverssm_ple.py +++ b/tests/kernels/mamba/test_recoverssm_ple.py @@ -124,7 +124,7 @@ def test_ple_commit_align_boundary(num_computed, accepted, crosses): mamba_block_size=block_size, ) - final_block = nxt if num_computed + accepted >= block_size else src + final_block = nxt if num_computed + accepted > block_size else src assert torch.equal(view[final_block, :, :STATE_LEN], want_final) if crosses: assert torch.equal(view[src, :, :STATE_LEN], want_boundary) diff --git a/tests/models/kimi_k3/test_recoverssm_boundary.py b/tests/models/kimi_k3/test_recoverssm_boundary.py new file mode 100644 index 0000000000..8023e44243 --- /dev/null +++ b/tests/models/kimi_k3/test_recoverssm_boundary.py @@ -0,0 +1,129 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Align-mode RecoverSSM state columns at recurrent block boundaries. + +A speculative window whose accepted tokens end exactly on a recurrent block +boundary must keep its state in the block holding the last committed token. +The following block can still be unallocated, so selecting it would drop the +commit while the bookkeeping recorded the state as advanced. +""" + +import pytest +import torch + +from vllm.models.kimi_k3.nvidia.ops.recoverssm import _prepare_commit_plan_kernel +from vllm.platforms import current_platform +from vllm.v1.attention.backends.utils import NULL_BLOCK_ID +from vllm.v1.worker.gpu.model_states.recoverssm import ( + _postprocess_recoverssm_align_kernel, +) + +pytestmark = pytest.mark.skipif( + not current_platform.is_cuda_alike(), + reason="The RecoverSSM kernels require a CUDA-alike device.", +) + +BLOCK_SIZE = 2048 +SPEC_QUERY_LEN = 4 +BLOCKS = [100, 101, 102, 103, 104, 105] + +# (computed before the step, committed tokens, allocated block columns, +# final column, boundary column or None) +CASES = [ + # 8,188 + 4 ends exactly on the 8,192 boundary; column 4 is unallocated. + (8188, 4, 4, 3, 3), + # The window crosses the boundary: state moves on, checkpoint stays behind. + (8190, 4, 5, 4, 3), + # The window stays inside column 3. + (8180, 4, 4, 3, None), + # The first token after a boundary belongs to the new column. + (8192, 1, 5, 4, None), + # Single-token commit ending on the first boundary. + (2047, 1, 1, 0, 0), +] + + +@pytest.mark.parametrize( + "num_computed,commit_len,allocated,final_col,boundary_col", CASES +) +def test_commit_plan_keeps_boundary_state_in_last_computed_block( + num_computed, commit_len, allocated, final_col, boundary_col +): + device = current_platform.device_type + columns = [ + block if col < allocated else NULL_BLOCK_ID for col, block in enumerate(BLOCKS) + ] + block_table = torch.tensor([columns], dtype=torch.int32, device=device) + num_accepted = torch.tensor([commit_len], dtype=torch.int32, device=device) + state_indices = torch.tensor( + [BLOCKS[(num_computed - 1) // BLOCK_SIZE]], dtype=torch.int32, device=device + ) + query_start_loc = torch.tensor( + [0, SPEC_QUERY_LEN], dtype=torch.int32, device=device + ) + computed = torch.tensor([num_computed], dtype=torch.int32, device=device) + commit_lens, final_indices, boundary_indices, recovery_lens = ( + torch.full((1,), -7, dtype=torch.int32, device=device) for _ in range(4) + ) + + _prepare_commit_plan_kernel[(1,)]( + num_accepted, + None, + state_indices, + query_start_loc, + block_table, + computed, + commit_lens, + final_indices, + boundary_indices, + recovery_lens, + NULL_BLOCK_ID, + BLOCK_SIZE, + block_table.shape[1], + num_accepted.stride(0), + 0, + state_indices.stride(0), + query_start_loc.stride(0), + block_table.stride(0), + block_table.stride(1), + computed.stride(0), + SPEC_QUERY_LEN=SPEC_QUERY_LEN, + num_warps=1, + ) + + assert commit_lens.item() == commit_len + assert final_indices.item() == BLOCKS[final_col] + expected_boundary = NULL_BLOCK_ID if boundary_col is None else BLOCKS[boundary_col] + assert boundary_indices.item() == expected_boundary + if boundary_col is not None: + next_boundary = (num_computed // BLOCK_SIZE + 1) * BLOCK_SIZE + assert recovery_lens.item() == next_boundary - num_computed + + +@pytest.mark.parametrize( + "num_computed,commit_len,allocated,final_col,boundary_col", CASES +) +def test_postprocess_records_the_committed_state_column( + num_computed, commit_len, allocated, final_col, boundary_col +): + device = current_platform.device_type + idx_mapping = torch.tensor([2], dtype=torch.int32, device=device) + num_sampled = torch.tensor([commit_len], dtype=torch.int32, device=device) + computed = torch.tensor([num_computed], dtype=torch.int32, device=device) + state_idx = torch.full((4,), -7, dtype=torch.int32, device=device) + num_accepted = torch.full((4,), -7, dtype=torch.int32, device=device) + + _postprocess_recoverssm_align_kernel[(1,)]( + idx_mapping, + num_sampled, + None, + computed, + state_idx, + num_accepted, + MAMBA_BLOCK_SIZE=BLOCK_SIZE, + BLOCK_TABLE_WIDTH=len(BLOCKS), + ) + + assert state_idx[2].item() == final_col + assert num_accepted[2].item() == 1 + assert state_idx[[0, 1, 3]].tolist() == [-7, -7, -7] diff --git a/vllm/model_executor/layers/mamba/ops/recoverssm.py b/vllm/model_executor/layers/mamba/ops/recoverssm.py index 073bdd55d4..9379688fcb 100644 --- a/vllm/model_executor/layers/mamba/ops/recoverssm.py +++ b/vllm/model_executor/layers/mamba/ops/recoverssm.py @@ -78,8 +78,12 @@ def prepare_commit_plan_kernel( tl.int32 ) final_num_computed = num_computed + commit_len + # The committed state belongs to the block holding the last committed + # token. A window ending exactly on a boundary stays in that block; the + # next block may still be unallocated. final_state_col = tl.minimum( - final_num_computed // mamba_block_size, block_table_width - 1 + tl.maximum(final_num_computed - 1, 0) // mamba_block_size, + block_table_width - 1, ) final_state_idx = tl.load( block_table_ptr diff --git a/vllm/v1/worker/gpu/model_states/recoverssm.py b/vllm/v1/worker/gpu/model_states/recoverssm.py index cfbe1a4d45..b02a6c19aa 100644 --- a/vllm/v1/worker/gpu/model_states/recoverssm.py +++ b/vllm/v1/worker/gpu/model_states/recoverssm.py @@ -91,10 +91,12 @@ def _postprocess_recoverssm_align_kernel( return num_sampled = tl.load(num_sampled_ptr + batch_idx) num_computed = tl.load(num_computed_ptr + batch_idx) + # Match the commit plan: the running state lives in the block holding the + # last computed token, so the next step migrates it across a boundary. tl.store( state_idx_ptr + req_state_idx, tl.minimum( - (num_computed + num_sampled) // MAMBA_BLOCK_SIZE, + tl.maximum(num_computed + num_sampled - 1, 0) // MAMBA_BLOCK_SIZE, BLOCK_TABLE_WIDTH - 1, ), )