[Bugfix] Route every speculation-capable row through the speculative path so recurrent state stays consistent (#56531)

Signed-off-by: Vadim Gimpelson <vadim.gimpelson@gmail.com>
Co-authored-by: Cursor Agent <cursoragent@cursor.com>
Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com>

tests/models/kimi_k3/test_kda_metadata.py

@@ -447,7 +447,7 @@ def test_internal_checkpoint_metadata_skips_unaligned_offset():
             2,
             False,
             [False, False],
-            None,
+            1,
             id="no-scheduled-draft-tokens",
         ),
         pytest.param(

tests/v1/attention/test_gdn_zero_draft_spec_rows.py

@@ -0,0 +1,211 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Which path GDN metadata takes when a decode row carries no draft tokens.
+
+`build_for_cudagraph_capture` derives the draft count from the query length
+(`num_decode_draft_tokens = query_len - 1`), so a capture shape of one token
+per request describes rows with zero drafts. `GDNAttentionMetadataBuilder`
+routes every row with `num_decode_draft_tokens_cpu >= 0` through the
+speculative path, zero drafts included: that path is the only one that
+applies each row's accepted-token offset to the recurrent state, so a
+zero-draft decode row must not fall back to the plain decode path.
+"""
+
+import pytest
+import torch
+
+from tests.v1.attention.utils import (
+    BatchSpec,
+    create_common_attn_metadata,
+    create_vllm_config,
+)
+from vllm.config import SpeculativeConfig
+from vllm.config.compilation import CUDAGraphMode
+from vllm.model_executor.layers.mamba.mamba_utils import MambaStateShapeCalculator
+from vllm.v1.attention.backends.gdn_attn import GDNAttentionMetadataBuilder
+from vllm.v1.kv_cache_interface import MambaSpec
+
+BLOCK_SIZE = 16
+DEVICE = torch.device("cpu")
+NUM_SPEC = 3
+
+
+def _builder(
+    use_spec_decode: bool,
+    mamba_cache_mode: str = "none",
+    shapes: tuple[tuple[int, ...], ...] = ((16, 64),),
+    dtypes: tuple[torch.dtype, ...] = (torch.float16,),
+) -> GDNAttentionMetadataBuilder:
+    vllm_config = create_vllm_config(
+        model_name="Qwen/Qwen3.5-0.8B",
+        block_size=BLOCK_SIZE,
+    )
+    vllm_config.compilation_config.cudagraph_mode = CUDAGraphMode.FULL_AND_PIECEWISE
+    vllm_config.speculative_config = (
+        SpeculativeConfig(method="ngram", num_speculative_tokens=NUM_SPEC)
+        if use_spec_decode
+        else None
+    )
+    # "none" keeps the block table as-is; the align gather needs columns for
+    # the speculative blocks, which `_decode_metadata` appends on request.
+    vllm_config.cache_config.mamba_cache_mode = mamba_cache_mode
+    mamba_spec = MambaSpec(
+        block_size=BLOCK_SIZE,
+        shapes=shapes,
+        dtypes=dtypes,
+        num_speculative_blocks=NUM_SPEC,
+    )
+
+    return GDNAttentionMetadataBuilder(
+        kv_cache_spec=mamba_spec,
+        layer_names=["layer.0"],
+        vllm_config=vllm_config,
+        device=DEVICE,
+    )
+
+
+def _capture_metadata(builder, query_len: int, batch_size: int = 2):
+    """Metadata as produced during CUDA-graph capture of a decode shape."""
+    batch = BatchSpec(
+        seq_lens=[64] * batch_size,
+        query_lens=[query_len] * batch_size,
+    )
+    common = create_common_attn_metadata(batch, BLOCK_SIZE, DEVICE)
+    return builder.build_for_cudagraph_capture(common)
+
+
+def _decode_metadata(builder, num_accepted: list[int], seq_len: int = 64):
+    """Metadata for a live decode batch where no row was given drafts.
+
+    Every row schedules a single token, so `num_decode_draft_tokens_cpu` is all
+    zeros -- the shape the runner produces once `zero_draft_decode_mask` has
+    promoted zero-draft decode rows out of the `-1` placeholder.
+    `num_accepted` is what the previous step accepted (bonus token included,
+    so one is the minimum), which is what decides the state slot.
+    """
+    batch = BatchSpec(
+        seq_lens=[seq_len] * len(num_accepted), query_lens=[1] * len(num_accepted)
+    )
+    common = create_common_attn_metadata(
+        batch, BLOCK_SIZE, DEVICE, arange_block_indices=True
+    )
+    if builder.vllm_config.cache_config.mamba_cache_mode == "align":
+        # Align gathers the last 1 + num_speculative_blocks columns starting at
+        # the request's current block, so the table has to reach past the end of
+        # the sequence -- at runtime it is sized by max_model_len, while the
+        # test helper sizes it by seq_len.
+        rows, cols = common.block_table_tensor.shape
+        extra = (
+            torch.arange(
+                rows * NUM_SPEC, dtype=common.block_table_tensor.dtype, device=DEVICE
+            ).view(rows, NUM_SPEC)
+            + rows * cols
+        )
+        common.block_table_tensor = torch.cat([common.block_table_tensor, extra], dim=1)
+    accepted = torch.tensor(num_accepted, dtype=torch.int32, device=DEVICE)
+    num_decode_draft_tokens_cpu = torch.zeros(len(num_accepted), dtype=torch.int32)
+    return builder.build(0, common, accepted, num_decode_draft_tokens_cpu)
+
+
+def test_zero_draft_capture_shape_routes_through_the_speculative_path():
+    """The one-token capture shape is what breaks backends whose decode kernels
+    assume a full num_spec + 1 query window: `num_decode_draft_tokens_cpu =
+    query_len - 1 = 0` for every row, which is `>= 0` and must still take the
+    speculative path (not the plain decode path that drops the state offset),
+    with a state tensor sized for the full window even though every row only
+    supplies a single query token.
+
+    Only `query_len=1` (zero drafts) discriminates this from the plain-decode
+    path before the fix: for any `query_len > 1`, `num_decode_draft_tokens_cpu`
+    is already `> 0` and both the pre-fix and post-fix code agree, so this
+    does not need parametrizing over larger draft counts.
+    """
+    meta = _capture_metadata(_builder(use_spec_decode=True), query_len=1)
+
+    assert meta.num_spec_decodes == 2
+    assert meta.num_decodes == 0
+    slots = meta.spec_state_indices_tensor
+    assert slots is not None
+    # The kernels index columns 0..num_spec of this tensor.
+    assert slots.shape[1] == NUM_SPEC + 1
+    assert int(slots.min()) >= 0
+    query_lens = torch.diff(meta.spec_query_start_loc)
+    assert torch.all(query_lens == 1)
+
+
+def test_plain_decode_without_spec_decode_skips_the_speculative_path():
+    """Without a speculative config the builder never builds spec metadata,
+    so a one-token-per-row decode capture takes the plain decode path.
+
+    Not a regression check for this patch (this branch is untouched by it),
+    but it is what keeps the test above from passing vacuously: it would also
+    pass if the builder routed every row through the speculative path
+    unconditionally, regardless of whether speculative decoding is even
+    configured.
+    """
+    meta = _capture_metadata(_builder(use_spec_decode=False), query_len=1)
+
+    assert meta.num_spec_decodes == 0
+    assert meta.num_decodes == 2
+
+
+@pytest.mark.parametrize(
+    "mamba_cache_mode, shapes, dtypes",
+    [
+        ("none", ((16, 64),), (torch.float16,)),
+        # Kimi-Linear's KDA layers subclass GatedDeltaNetAttention and consume
+        # this same metadata, and prefix caching (which Kimi enables by
+        # default) puts the cache in align mode -- so the guarantee has to hold
+        # for that state shape and mode too.
+        (
+            "align",
+            MambaStateShapeCalculator.kda_state_shape(
+                tp_world_size=1, num_heads=4, head_dim=32, num_spec=NUM_SPEC
+            ),
+            (torch.bfloat16, torch.bfloat16),
+        ),
+    ],
+    ids=["gdn_none", "kimi_kda_align"],
+)
+def test_zero_draft_batch_still_applies_the_accepted_token_offset(
+    mamba_cache_mode, shapes, dtypes
+):
+    """The numeric difference the fix makes, on the batch shape that triggers it.
+
+    The recurrent kernels read their state at column `num_accepted_tokens - 1`
+    of `spec_state_indices_tensor`. Before the fix, a batch in which *every*
+    speculative row drafted nothing was pushed onto the non-speculative path,
+    which passes `non_spec_state_indices_tensor` -- always column 0 -- and
+    drops `num_accepted_tokens` entirely. A row that accepted drafts on the
+    previous step wrote its fresh state further along, so column 0 is stale.
+    """
+    accepted = [NUM_SPEC, 1]
+    meta = _decode_metadata(
+        _builder(
+            use_spec_decode=True,
+            mamba_cache_mode=mamba_cache_mode,
+            shapes=shapes,
+            dtypes=dtypes,
+        ),
+        accepted,
+    )
+
+    assert meta.num_spec_decodes == len(accepted)
+    assert meta.num_decodes == 0
+    # Dropping these is what silently sent the kernels to the wrong slot.
+    assert meta.num_accepted_tokens is not None
+    assert meta.num_accepted_tokens.tolist()[: len(accepted)] == accepted
+    assert meta.non_spec_state_indices_tensor is None
+
+    slots = meta.spec_state_indices_tensor
+    assert slots is not None
+    for row, num_accepted in enumerate(accepted):
+        read = int(slots[row, num_accepted - 1])
+        fallback = int(slots[row, 0])
+        # Only the row that accepted a draft moved off column 0; the row that
+        # accepted nothing but its bonus token legitimately stays there. This
+        # also confirms the columns are genuinely distinct blocks (via
+        # `arange_block_indices=True` in `_decode_metadata`), not just
+        # distinct labels for the same one -- otherwise `read != fallback`
+        # here would mean nothing.
+        assert (read != fallback) == (num_accepted > 1)

tests/v1/worker/test_gpu_model_runner_prefill_tail_classification.py

@@ -0,0 +1,249 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""`GPUModelRunner._prepare_inputs` must tell a chunked-prefill continuation
+row apart from a genuine zero-draft decode row.
+
+Both kinds of row are absent from `scheduled_spec_decode_tokens` and both
+schedule exactly one token, so `num_decode_draft_tokens` starts at the
+placeholder `-1` for either. The `zero_draft_decode_mask` in
+`gpu_model_runner.py` is what tells them apart afterwards, using
+`num_computed_tokens >= num_prompt_tokens`: a prefill continuation must keep
+`-1` (it carries no drafts and no accepted-token offset to apply), while a
+zero-draft decode row must become `0` (it still has to run the speculative
+path so the recurrent backends apply that offset). Getting this wrong in
+either direction was the root cause behind issue #40875: Kimi-Linear crashed
+with an illegal memory access when prefill tails were misclassified as
+zero-draft decodes.
+
+Each batch below also includes one row with a real draft token, so
+`scheduled_spec_decode_tokens` is non-empty even under the pre-fix
+`use_spec_decode = len(scheduled_spec_decode_tokens) > 0` predicate. That
+keeps the assertions targeted at the mask fix alone, not the unrelated
+`use_spec_decode` predicate fix (`self.speculative_config is not None`) that
+landed in the same patch.
+"""
+
+import numpy as np
+import pytest
+import torch
+
+from vllm.config import (
+    CacheConfig,
+    ModelConfig,
+    ParallelConfig,
+    SchedulerConfig,
+    SpeculativeConfig,
+    VllmConfig,
+    set_current_vllm_config,
+)
+from vllm.model_executor.layers.attention import Attention
+from vllm.sampling_params import SamplingParams
+from vllm.v1.core.sched.output import CachedRequestData, NewRequestData, SchedulerOutput
+from vllm.v1.kv_cache_interface import (
+    FullAttentionSpec,
+    KVCacheConfig,
+    KVCacheGroupSpec,
+    KVCacheTensor,
+)
+from vllm.v1.worker.gpu_input_batch import InputBatch
+from vllm.v1.worker.gpu_model_runner import GPUModelRunner
+
+BLOCK_SIZE = 16
+NUM_BLOCKS = 10
+DEVICE_TYPE = "cuda" if torch.cuda.is_available() else "cpu"
+
+
+def _get_vllm_config() -> VllmConfig:
+    model_config = ModelConfig(model="facebook/opt-125m", dtype="float16", seed=42)
+    scheduler_config = SchedulerConfig(
+        max_num_seqs=10,
+        max_num_batched_tokens=512,
+        max_model_len=512,
+        is_encoder_decoder=model_config.is_encoder_decoder,
+    )
+    cache_config = CacheConfig(
+        block_size=BLOCK_SIZE, gpu_memory_utilization=0.9, cache_dtype="auto"
+    )
+    cache_config.kv_cache_layout = "LBNHC"
+    return VllmConfig(
+        model_config=model_config,
+        cache_config=cache_config,
+        scheduler_config=scheduler_config,
+        parallel_config=ParallelConfig(),
+        # ngram needs no draft model weights, so the fixture stays as cheap as
+        # the speculative-decoding-free one in test_gpu_model_runner.py.
+        speculative_config=SpeculativeConfig(method="ngram", num_speculative_tokens=3),
+    )
+
+
+def _initialize_kv_cache(runner: GPUModelRunner) -> None:
+    attn_spec = FullAttentionSpec(
+        block_size=BLOCK_SIZE,
+        num_kv_heads=runner.model_config.get_num_kv_heads(runner.parallel_config),
+        head_size=runner.model_config.get_head_size(),
+        dtype=runner.kv_cache_dtype,
+    )
+    tensor_size = attn_spec.page_size_bytes * NUM_BLOCKS
+    kv_cache_config = KVCacheConfig(
+        num_blocks=NUM_BLOCKS,
+        kv_cache_tensors=[
+            KVCacheTensor(
+                size=tensor_size,
+                layers=["layer.0"],
+                layer_stride=tensor_size,
+                block_stride=attn_spec.page_size_bytes,
+            ),
+        ],
+        kv_cache_groups=[
+            KVCacheGroupSpec(layer_names=["layer.0"], kv_cache_spec=attn_spec)
+        ],
+    )
+    runner.kv_cache_config = kv_cache_config
+    runner.input_batch = InputBatch(
+        max_num_reqs=runner.max_num_reqs,
+        max_model_len=runner.max_model_len,
+        max_num_batched_tokens=runner.max_num_tokens,
+        device=runner.device,
+        vocab_size=runner.model_config.get_vocab_size(),
+        block_sizes=[kv_cache_config.kv_cache_groups[0].kv_cache_spec.block_size],
+        kernel_block_sizes=[
+            kv_cache_config.kv_cache_groups[0].kv_cache_spec.block_size
+        ],
+        max_num_blocks_per_req=[NUM_BLOCKS],
+    )
+    runner.initialize_attn_backend(kv_cache_config)
+
+
+@pytest.fixture
+def spec_decode_model_runner(dist_init):
+    vllm_config = _get_vllm_config()
+    with set_current_vllm_config(vllm_config):
+        model_config = vllm_config.model_config
+        num_heads = model_config.get_num_kv_heads(vllm_config.parallel_config)
+        head_size = model_config.get_head_size()
+        vllm_config.compilation_config.static_forward_context["layer.0"] = Attention(
+            num_heads, head_size, 0.1
+        )
+        runner = GPUModelRunner(vllm_config, DEVICE_TYPE)
+        _initialize_kv_cache(runner)
+        yield runner
+
+
+def _new_request(
+    req_id: str, prompt_len: int, num_computed_tokens: int
+) -> tuple[NewRequestData, int]:
+    """A request with `prompt_len` prompt tokens, `num_computed_tokens` of
+    which are already computed (e.g. a prefix-cache hit), scheduled for the
+    single remaining/next token this step."""
+    return (
+        NewRequestData(
+            req_id=req_id,
+            prompt_token_ids=list(range(prompt_len)),
+            mm_features=[],
+            sampling_params=SamplingParams(),
+            pooling_params=None,
+            block_ids=([0],),
+            num_computed_tokens=num_computed_tokens,
+            lora_request=None,
+        ),
+        1,
+    )
+
+
+def test_prefill_continuation_keeps_placeholder_decode_zero_draft_becomes_zero(
+    spec_decode_model_runner,
+):
+    """One request has one token left to prefill (`num_computed_tokens ==
+    num_prompt_tokens - 1`); another has already finished its prompt and is
+    now decoding with no draft tokens proposed this step; a third is decoding
+    with a real draft. `decode_with_draft` gives the batch a non-empty
+    `scheduled_spec_decode_tokens`, so this also exercises the `-1`-vs-`0`
+    distinction on the pre-fix definition of `use_spec_decode` (`len(sched
+    uled_spec_decode_tokens) > 0`), isolating the mask fix from the unrelated
+    `use_spec_decode` predicate fix that landed in the same patch."""
+    runner = spec_decode_model_runner
+
+    prefill_req, prefill_sched = _new_request(
+        "prefill_tail", prompt_len=11, num_computed_tokens=10
+    )
+    decode_req, decode_sched = _new_request(
+        "decode_zero_draft", prompt_len=3, num_computed_tokens=3
+    )
+    drafted_req, _ = _new_request(
+        "decode_with_draft", prompt_len=3, num_computed_tokens=3
+    )
+    num_scheduled_tokens = {
+        "prefill_tail": prefill_sched,
+        "decode_zero_draft": decode_sched,
+        "decode_with_draft": 2,  # 1 draft token + 1 bonus token
+    }
+    scheduler_output = SchedulerOutput(
+        scheduled_new_reqs=[prefill_req, decode_req, drafted_req],
+        scheduled_cached_reqs=CachedRequestData.make_empty(),
+        num_scheduled_tokens=num_scheduled_tokens,
+        total_num_scheduled_tokens=sum(num_scheduled_tokens.values()),
+        scheduled_spec_decode_tokens={"decode_with_draft": [99]},
+        scheduled_encoder_inputs={},
+        num_common_prefix_blocks=[],
+        finished_req_ids=set(),
+        free_encoder_mm_hashes=[],
+    )
+
+    runner._update_states(scheduler_output)
+    runner._prepare_inputs(
+        scheduler_output,
+        np.array(
+            [num_scheduled_tokens[r] for r in runner.input_batch.req_ids],
+            dtype=np.int32,
+        ),
+    )
+
+    prefill_idx = runner.input_batch.req_id_to_index["prefill_tail"]
+    decode_idx = runner.input_batch.req_id_to_index["decode_zero_draft"]
+    drafted_idx = runner.input_batch.req_id_to_index["decode_with_draft"]
+    assert runner.num_decode_draft_tokens.np[prefill_idx] == -1
+    assert runner.num_decode_draft_tokens.np[decode_idx] == 0
+    assert runner.num_decode_draft_tokens.np[drafted_idx] == 1
+
+
+def test_prefill_tail_of_two_or_three_tokens_also_keeps_the_placeholder(
+    spec_decode_model_runner,
+):
+    """A chunked-prefill tail can be 1, 2, or 3 tokens (`long_prefill_token_
+    threshold` rounding); the mask only fires for `num_scheduled_tokens == 1`,
+    so wider tails must keep `-1` regardless of how much prompt is left. A
+    third, drafted decode row keeps `scheduled_spec_decode_tokens` non-empty
+    so the assertions hold on the pre-fix `use_spec_decode` predicate too."""
+    runner = spec_decode_model_runner
+
+    tail_2_req, _ = _new_request("tail_2", prompt_len=12, num_computed_tokens=10)
+    tail_3_req, _ = _new_request("tail_3", prompt_len=13, num_computed_tokens=10)
+    drafted_req, _ = _new_request(
+        "decode_with_draft", prompt_len=3, num_computed_tokens=3
+    )
+    num_scheduled_tokens = {"tail_2": 2, "tail_3": 3, "decode_with_draft": 2}
+    scheduler_output = SchedulerOutput(
+        scheduled_new_reqs=[tail_2_req, tail_3_req, drafted_req],
+        scheduled_cached_reqs=CachedRequestData.make_empty(),
+        num_scheduled_tokens=num_scheduled_tokens,
+        total_num_scheduled_tokens=sum(num_scheduled_tokens.values()),
+        scheduled_spec_decode_tokens={"decode_with_draft": [99]},
+        scheduled_encoder_inputs={},
+        num_common_prefix_blocks=[],
+        finished_req_ids=set(),
+        free_encoder_mm_hashes=[],
+    )
+
+    runner._update_states(scheduler_output)
+    runner._prepare_inputs(
+        scheduler_output,
+        np.array(
+            [num_scheduled_tokens[r] for r in runner.input_batch.req_ids],
+            dtype=np.int32,
+        ),
+    )
+
+    tail_2_idx = runner.input_batch.req_id_to_index["tail_2"]
+    tail_3_idx = runner.input_batch.req_id_to_index["tail_3"]
+    assert runner.num_decode_draft_tokens.np[tail_2_idx] == -1
+    assert runner.num_decode_draft_tokens.np[tail_3_idx] == -1

tests/v1/worker/test_mamba_hybrid_model_state.py

@@ -140,6 +140,177 @@ def test_padded_prompt_tail_builds_as_spec_decode(
     assert meta.num_prefill_tokens == k + 1
 
 
+def _mamba_hybrid_state(num_speculative_tokens: int) -> MambaHybridModelState:
+    state = object.__new__(MambaHybridModelState)
+    state.vllm_config = SimpleNamespace(num_speculative_tokens=num_speculative_tokens)
+    state.model_config = SimpleNamespace(max_model_len=8192)
+    state._align_mode = False
+    state.recoverssm = None
+    # A request that accepted 3 drafts (plus the bonus token) on the previous
+    # step; the others have never had a draft accepted.
+    state.num_accepted_tokens_gpu = torch.tensor([4, 1, 1], dtype=torch.int32)
+    return state
+
+
+def _input_batch(
+    num_scheduled_tokens: list[int],
+    num_draft_tokens_per_req: list[int] | None,
+    is_prefilling: list[bool],
+    num_reqs_after_padding: int | None = None,
+) -> SimpleNamespace:
+    """`num_draft_tokens_per_req=None` is a batch where nobody proposed a
+    draft this step, e.g. because ngram found no match for anyone."""
+    num_reqs = len(num_scheduled_tokens)
+    query_start_loc = np.zeros(num_reqs + 1, dtype=np.int32)
+    query_start_loc[1:] = np.cumsum(num_scheduled_tokens)
+    num_tokens = int(query_start_loc[-1])
+    seq_lens = torch.full((num_reqs,), 10, dtype=torch.int32)
+    return SimpleNamespace(
+        num_reqs=num_reqs,
+        num_reqs_after_padding=num_reqs_after_padding or num_reqs,
+        num_tokens=num_tokens,
+        num_tokens_after_padding=num_tokens,
+        query_start_loc_np=query_start_loc,
+        num_scheduled_tokens=np.array(num_scheduled_tokens, dtype=np.int32),
+        max_query_len=None,
+        seq_lens_cpu_upper_bound=seq_lens,
+        is_prefilling_np=np.array(is_prefilling),
+        prefill_runs_as_decode_np=None,
+        num_draft_tokens_per_req=None
+        if num_draft_tokens_per_req is None
+        else np.array(num_draft_tokens_per_req, dtype=np.int32),
+        idx_mapping=torch.arange(num_reqs),
+        query_start_loc=torch.from_numpy(query_start_loc),
+        seq_lens=seq_lens,
+        dcp_local_seq_lens=None,
+        positions=torch.zeros(num_tokens, dtype=torch.int64),
+        prompt_lens=None,
+    )
+
+
+def _prepare_attn_metadata(
+    state: MambaHybridModelState,
+    monkeypatch,
+    input_batch: SimpleNamespace | None = None,
+    cudagraph_mode: CUDAGraphMode = CUDAGraphMode.NONE,
+) -> SimpleNamespace:
+    captured = {}
+
+    def fake_build_attn_metadata(**kwargs):
+        captured.update(kwargs)
+        return Mock()
+
+    monkeypatch.setattr(
+        "vllm.v1.worker.gpu.model_states.mamba_hybrid.build_attn_metadata",
+        fake_build_attn_metadata,
+    )
+
+    if input_batch is None:
+        input_batch = _input_batch([1, 1], None, [False, False])
+    state.prepare_attn(
+        input_batch=input_batch,
+        cudagraph_mode=cudagraph_mode,
+        block_tables=(),
+        slot_mappings={},
+        attn_groups=[],
+        kv_cache_config=Mock(),
+    )
+    return captured["model_specific_attn_metadata"]
+
+
+def test_prepare_attn_keeps_the_accepted_offset_when_the_batch_drafted_nothing(
+    monkeypatch,
+):
+    """`num_accepted_tokens` must reach the attention builders whenever the
+    speculative config is on, even when *this* batch scheduled no draft
+    tokens at all. Before the fix, `use_spec_decode` was gated on this step's
+    scheduled drafts (`len(scheduled_spec_decode_tokens) > 0` in the V1
+    runner) rather than on the speculative config, so a request that accepted
+    drafts on a previous step would silently lose that offset the moment
+    ngram found no match for anyone in the batch.
+    """
+    state = _mamba_hybrid_state(num_speculative_tokens=3)
+
+    metadata = _prepare_attn_metadata(state, monkeypatch)
+
+    assert metadata.num_accepted_tokens is not None
+    assert metadata.num_accepted_tokens.tolist() == [4, 1]
+    assert metadata.num_decode_draft_tokens_cpu is not None
+
+
+def test_prepare_attn_omits_the_offset_without_a_speculative_config(monkeypatch):
+    """Mirrors the check above: with speculative decoding off entirely, there
+    is no accepted-token offset to carry, and the metadata must say so
+    explicitly (`None`) rather than a stale or default value."""
+    state = _mamba_hybrid_state(num_speculative_tokens=0)
+
+    metadata = _prepare_attn_metadata(state, monkeypatch)
+
+    assert metadata.num_accepted_tokens is None
+    assert metadata.num_decode_draft_tokens_cpu is None
+
+
+def test_zero_draft_decode_row_is_marked_as_a_speculative_row(monkeypatch):
+    """A decode row that got no drafts still has to reach the speculative path.
+
+    The recurrent backends only apply a row's accepted-token offset on that
+    path, and a row that accepted drafts on the previous step keeps that offset
+    even when the current step proposes nothing (the scheduler drops drafts
+    whenever the token budget truncates the step). Marking such a row with the
+    -1 sentinel sends it to the plain decode path, which reads the state at
+    column 0 and ignores the offset.
+    """
+    state = _mamba_hybrid_state(num_speculative_tokens=3)
+    input_batch = _input_batch([3, 1, 1], [2, 0, 0], [False, False, True])
+
+    metadata = _prepare_attn_metadata(state, monkeypatch, input_batch)
+
+    assert metadata.num_decode_draft_tokens_cpu.tolist() == [
+        2,  # decode row with drafts
+        0,  # decode row with no drafts: speculative path, offset applied
+        -1,  # one-token prefill tail: no offset to apply, stays a plain row
+    ]
+
+
+def test_batch_without_any_drafts_still_marks_its_decode_rows(monkeypatch):
+    """`num_draft_tokens_per_req` is None whenever no row in the batch drafted
+    anything. Those rows are ordinary decode rows carrying an offset from the
+    previous step, so the absence of drafts in this step must not be read as
+    speculative decoding being off."""
+    state = _mamba_hybrid_state(num_speculative_tokens=3)
+    input_batch = _input_batch([1, 1], None, [False, True])
+
+    metadata = _prepare_attn_metadata(state, monkeypatch, input_batch)
+
+    assert metadata.num_decode_draft_tokens_cpu.tolist() == [0, -1]
+
+
+def test_padded_rows_keep_the_sentinel(monkeypatch):
+    """Full CUDA-graph capture pads the batch; padded rows describe no request
+    and must not be picked up as speculative rows."""
+    state = _mamba_hybrid_state(num_speculative_tokens=3)
+    input_batch = _input_batch([1, 1], [0, 0], [False, False], num_reqs_after_padding=4)
+
+    metadata = _prepare_attn_metadata(
+        state, monkeypatch, input_batch, cudagraph_mode=CUDAGraphMode.FULL
+    )
+
+    assert metadata.num_decode_draft_tokens_cpu.tolist() == [0, 0, -1, -1]
+
+
+def test_chunked_prefill_tail_of_two_or_three_tokens_keeps_the_sentinel(
+    monkeypatch,
+):
+    """A prefill continuation is excluded by `is_prefilling`, not by its token
+    count, so tails wider than one token are covered by the same rule."""
+    state = _mamba_hybrid_state(num_speculative_tokens=3)
+    input_batch = _input_batch([2, 3], [0, 0], [True, True])
+
+    metadata = _prepare_attn_metadata(state, monkeypatch, input_batch)
+
+    assert metadata.num_decode_draft_tokens_cpu.tolist() == [-1, -1]
+
+
 @pytest.mark.skipif(not current_platform.is_cuda(), reason="Requires CUDA")
 @pytest.mark.parametrize(("num_sampled", "expected_value"), [(0, 1), (3, 3)])
 def test_postprocess_state_scalar_with_int32_mapping(

vllm/models/kimi_k3/nvidia/kda_metadata.py

@@ -409,20 +409,12 @@ class KimiK3KDAMetadataBuilder(GDNAttentionMetadataBuilder):
                 spec_sequence_masks_cpu &= (
                     query_start_loc_cpu.diff() <= self.num_spec + 1
                 )
-            # Native KDA can use its regular decode path when no draft token
-            # was scheduled. RecoverSSM must preserve its extended conv window.
-            if (
-                not self.use_recoverssm
-                and num_decode_draft_tokens_cpu[spec_sequence_masks_cpu].sum().item()
-                == 0
-            ):
+            # A batch whose rows all drafted nothing still has to run the
+            # speculative path: that is the only path that applies each row's
+            # accepted-token offset to the recurrent state.
+            num_spec_decodes = spec_sequence_masks_cpu.sum().item()
+            if num_spec_decodes == 0:
                 spec_sequence_masks_cpu = None
-                num_spec_decodes = 0
-                active_non_spec_mask_cpu = None
-            else:
-                num_spec_decodes = spec_sequence_masks_cpu.sum().item()
-                if num_spec_decodes == 0:
-                    spec_sequence_masks_cpu = None
 
         spec_request_indices = None
         spec_token_start = None

vllm/v1/attention/backends/gdn_attn.py

@@ -260,11 +260,10 @@ class GDNAttentionMetadataBuilder(AttentionMetadataBuilder[GDNAttentionMetadata]
         else:
             spec_sequence_masks_cpu = num_decode_draft_tokens_cpu >= 0
             num_spec_decodes = spec_sequence_masks_cpu.sum().item()
-            if (
-                num_spec_decodes == 0
-                or num_decode_draft_tokens_cpu[spec_sequence_masks_cpu].sum().item()
-                == 0
-            ):
+            # A batch whose rows all drafted nothing still has to run the
+            # speculative path: that is the only path that applies each row's
+            # accepted-token offset to the recurrent state.
+            if num_spec_decodes == 0:
                 num_spec_decodes = 0
                 spec_sequence_masks = None
                 spec_sequence_masks_cpu = None

vllm/v1/worker/gpu/model_states/mamba_hybrid.py

@@ -277,18 +277,19 @@ class MambaHybridModelState(DefaultModelState):
             ]
 
             # GDN uses >= 0 to select spec-decode rows, so non-decode rows
-            # need the -1 sentinel rather than a raw zero draft count.
+            # need the -1 sentinel. Decode rows without drafts stay spec rows:
+            # only that path applies the previous step's accepted offset.
             num_decode_draft_tokens_np = np.full(num_reqs, -1, dtype=np.int32)
             num_draft_tokens_per_req = input_batch.num_draft_tokens_per_req
-            if num_draft_tokens_per_req is not None:
-                # Test request state, not num_scheduled_tokens == draft_count+1:
-                # adaptive rewrites num_scheduled_tokens to an even split, so that
-                # equality rarely holds and would demote every verify row to decode.
-                is_decode = ~is_prefilling_np & (input_batch.num_scheduled_tokens > 0)
-                spec_decode_mask = (num_draft_tokens_per_req > 0) & is_decode
-                num_decode_draft_tokens_np[: input_batch.num_reqs] = np.where(
-                    spec_decode_mask, num_draft_tokens_per_req, -1
-                )
+            if num_draft_tokens_per_req is None:
+                num_draft_tokens_per_req = np.zeros(input_batch.num_reqs, np.int32)
+            # Test request state, not num_scheduled_tokens == draft_count+1:
+            # adaptive rewrites num_scheduled_tokens to an even split, so that
+            # equality rarely holds and would demote every verify row to decode.
+            is_decode = ~is_prefilling_np & (input_batch.num_scheduled_tokens > 0)
+            num_decode_draft_tokens_np[: input_batch.num_reqs] = np.where(
+                is_decode, num_draft_tokens_per_req, -1
+            )
             num_decode_draft_tokens_cpu = torch.from_numpy(num_decode_draft_tokens_np)
 
         if self._align_mode:

vllm/v1/worker/gpu_model_runner.py

@@ -840,6 +840,12 @@ class GPUModelRunner(
         self.num_accepted_tokens = self._make_buffer(
             self.max_num_reqs, dtype=torch.int32
         )
+        # The count includes the bonus token, so one is the minimum, and the
+        # recurrent kernels pick their state slot as `accepted - 1`. Dummy runs
+        # skip `_prepare_inputs`, so the zeros from the allocation would reach
+        # those kernels as index -1 and fault on the out-of-bounds read.
+        self.num_accepted_tokens.np.fill(1)
+        self.num_accepted_tokens.copy_to_gpu()
 
         # Only relevant for models using M-RoPE (e.g, Qwen2-VL)
         if self.uses_mrope:
@@ -2199,7 +2205,7 @@ class GPUModelRunner(
             )
             self.mrope_positions.gpu[:, :total_num_scheduled_tokens] += drift
 
-        use_spec_decode = len(scheduler_output.scheduled_spec_decode_tokens) > 0
+        use_spec_decode = self.speculative_config is not None
         if not use_spec_decode:
             # NOTE(woosuk): Due to chunked prefills, the batch may contain
             # partial requests. While we should not sample any token
@@ -2226,6 +2232,18 @@ class GPUModelRunner(
                 num_draft_tokens[req_idx] = draft_len
                 if num_scheduled_tokens[req_idx] == draft_len + 1:
                     num_decode_draft_tokens[req_idx] = draft_len
+            # Some recurrent backends may read the previous step's state at an
+            # offset that only the speculative path applies. So a decode step
+            # that got no drafts must still count as a speculative row with
+            # zero drafts, not as a plain decode. Prompt chunks are left out:
+            # their tokens are prompt tokens and carry no such offset.
+            zero_draft_decode_mask = (
+                self.input_batch.num_computed_tokens_cpu[:num_reqs]
+                >= self.input_batch.num_prompt_tokens[:num_reqs]
+            ) & (num_scheduled_tokens[:num_reqs] == 1)
+            num_decode_draft_tokens[
+                (num_decode_draft_tokens < 0) & zero_draft_decode_mask
+            ] = 0
             spec_decode_metadata = self._calc_spec_decode_metadata(
                 num_draft_tokens, cu_num_tokens
             )
@@ -4312,7 +4330,7 @@ class GPUModelRunner(
                         self.mamba_state_idx,
                     )
 
-            use_spec_decode = len(scheduler_output.scheduled_spec_decode_tokens) > 0
+            use_spec_decode = self.speculative_config is not None
             ubatch_slices_attn = ubatch_slices_padded if pad_attn else ubatch_slices
 
             slot_mappings_by_group, slot_mappings = self._get_slot_mappings(