From 0eac15270710ddb7f5296d7244fcd0496c2daa0f Mon Sep 17 00:00:00 2001 From: Vadim Gimpelson <156319763+vadiklyutiy@users.noreply.github.com> Date: Tue, 6 Oct 2026 03:03:44 +0400 Subject: [PATCH] [Bugfix] Route every speculation-capable row through the speculative path so recurrent state stays consistent (#56531) Signed-off-by: Vadim Gimpelson Co-authored-by: Cursor Agent Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com> --- tests/models/kimi_k3/test_kda_metadata.py | 2 +- .../test_gdn_zero_draft_spec_rows.py | 211 +++++++++++++++ ...odel_runner_prefill_tail_classification.py | 249 ++++++++++++++++++ .../worker/test_mamba_hybrid_model_state.py | 171 ++++++++++++ vllm/models/kimi_k3/nvidia/kda_metadata.py | 18 +- vllm/v1/attention/backends/gdn_attn.py | 9 +- .../worker/gpu/model_states/mamba_hybrid.py | 21 +- vllm/v1/worker/gpu_model_runner.py | 22 +- 8 files changed, 672 insertions(+), 31 deletions(-) create mode 100644 tests/v1/attention/test_gdn_zero_draft_spec_rows.py create mode 100644 tests/v1/worker/test_gpu_model_runner_prefill_tail_classification.py diff --git a/tests/models/kimi_k3/test_kda_metadata.py b/tests/models/kimi_k3/test_kda_metadata.py index 6b3c71bdc6..7441427955 100644 --- a/tests/models/kimi_k3/test_kda_metadata.py +++ b/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( diff --git a/tests/v1/attention/test_gdn_zero_draft_spec_rows.py b/tests/v1/attention/test_gdn_zero_draft_spec_rows.py new file mode 100644 index 0000000000..55955886ce --- /dev/null +++ b/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) diff --git a/tests/v1/worker/test_gpu_model_runner_prefill_tail_classification.py b/tests/v1/worker/test_gpu_model_runner_prefill_tail_classification.py new file mode 100644 index 0000000000..9375c477a2 --- /dev/null +++ b/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 diff --git a/tests/v1/worker/test_mamba_hybrid_model_state.py b/tests/v1/worker/test_mamba_hybrid_model_state.py index dba5444150..f87201e48d 100644 --- a/tests/v1/worker/test_mamba_hybrid_model_state.py +++ b/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( diff --git a/vllm/models/kimi_k3/nvidia/kda_metadata.py b/vllm/models/kimi_k3/nvidia/kda_metadata.py index e5a63f13a3..58c74e6258 100644 --- a/vllm/models/kimi_k3/nvidia/kda_metadata.py +++ b/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 diff --git a/vllm/v1/attention/backends/gdn_attn.py b/vllm/v1/attention/backends/gdn_attn.py index 5fece32a1b..95f6a808a5 100644 --- a/vllm/v1/attention/backends/gdn_attn.py +++ b/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 diff --git a/vllm/v1/worker/gpu/model_states/mamba_hybrid.py b/vllm/v1/worker/gpu/model_states/mamba_hybrid.py index 32924555b1..f3f3c666e0 100644 --- a/vllm/v1/worker/gpu/model_states/mamba_hybrid.py +++ b/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: diff --git a/vllm/v1/worker/gpu_model_runner.py b/vllm/v1/worker/gpu_model_runner.py index df7dc96516..76299308f8 100644 --- a/vllm/v1/worker/gpu_model_runner.py +++ b/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(