Add vLLM v0.18.1 source tree with KV transfer abort fix

third_party/vllm/ now tracked in git for direct patch management.
Based on vLLM v0.18.1 release with one patch applied:

  vllm/v1/core/sched/scheduler.py:
    Replace fatal assert with graceful skip when KV transfer callback
    arrives for an already-aborted request during PD disaggregated serving.

Future vLLM modifications should be made directly in third_party/vllm/
and committed normally. The patches/ directory is kept as documentation
of what changed from upstream.

Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>

third_party/vllm/tests/v1/e2e/general/test_async_scheduling.py

@@ -0,0 +1,437 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+import os
+from itertools import repeat
+from typing import Any
+
+import pytest
+import torch._dynamo.config as dynamo_config
+
+from tests.utils import (
+    large_gpu_mark,
+    single_gpu_only,
+)
+from vllm import SamplingParams
+from vllm.logprobs import Logprob
+from vllm.platforms import current_platform
+from vllm.sampling_params import StructuredOutputsParams
+from vllm.v1.metrics.reader import Metric
+
+from ....conftest import VllmRunner
+from ....models.utils import check_outputs_equal
+
+MODEL = "Qwen/Qwen3-0.6B"
+MTP_MODEL = "meta-llama/Llama-3.2-1B-Instruct"
+
+# Need to enforce eager for MRV2 while we sort out cudagraph issues.
+ENFORCE_EAGER = os.getenv("ENFORCE_EAGER", "0") == "1"
+
+first_prompt = (
+    "The following numbers of the sequence "
+    + ", ".join(str(i) for i in range(10))
+    + " are:"
+)
+example_prompts = [first_prompt, "In one word, the capital of France is "] + [
+    f"Tell me about the number {i}: " for i in range(32)
+]
+
+default_params = dict(
+    temperature=0.0,  # greedy
+    max_tokens=30,
+    min_tokens=28,
+)
+
+
+@single_gpu_only
+def test_without_spec_decoding(
+    sample_json_schema,
+    monkeypatch: pytest.MonkeyPatch,
+):
+    """Test consistency of combos of async scheduling, preemption,
+    uni/multiproc executor, prefill chunking."""
+    struct_outputs = StructuredOutputsParams(json=sample_json_schema)
+    test_sampling_params: list[dict[str, Any]] = [
+        dict(),
+        # dict(min_tokens=20),
+        dict(frequency_penalty=-1.0),
+        dict(bad_words=["the", " the"]),
+        dict(logprobs=2),
+        dict(logprobs=2, frequency_penalty=-1.0),
+        dict(structured_outputs=struct_outputs),
+        dict(
+            structured_outputs=struct_outputs,
+            logprobs=2,
+        ),
+        dict(
+            structured_outputs=struct_outputs,
+            frequency_penalty=-1.0,
+        ),
+        dict(
+            structured_outputs=struct_outputs,
+            logprobs=2,
+            frequency_penalty=-1.0,
+        ),
+    ]
+
+    # test_preemption, executor, async_scheduling,
+    # spec_config, test_prefill_chunking
+    test_configs = [
+        (False, "mp", False, None, False),
+        (True, "mp", False, None, True),
+        (False, "mp", True, None, False),
+        (False, "uni", True, None, False),
+        (True, "mp", True, None, False),
+        (True, "uni", True, None, False),
+        (False, "mp", True, None, True),
+        (True, "mp", True, None, True),
+        (True, "uni", True, None, True),
+    ]
+
+    if current_platform.is_rocm():
+        # On ROCm, Only test with structured_outputs (deterministic)
+        # and skip chunk_prefill (more variable).
+        test_configs = [
+            cfg
+            for cfg in test_configs
+            if not cfg[4]  # skip chunk_prefill=True
+        ]
+        test_sampling_params = [
+            p for p in test_sampling_params if p.get("structured_outputs") is not None
+        ]
+
+    run_tests(monkeypatch, MODEL, test_configs, test_sampling_params)
+
+
+@single_gpu_only
+@large_gpu_mark(min_gb=16)
+def test_with_eagle3_spec_decoding(sample_json_schema, monkeypatch: pytest.MonkeyPatch):
+    """Test consistency and acceptance rates with some different combos of
+    preemption, executor, async scheduling, prefill chunking,
+    spec decoding model length.
+    """
+
+    spec_config = {
+        "method": "eagle3",
+        "num_speculative_tokens": 2,
+        "model": "nm-testing/Llama3_2_1B_speculator.eagle3",
+    }
+    # Set small draft model len to force doesn't-fit-in-drafter case.
+    spec_config_short = spec_config | {"max_model_len": 50}
+
+    struct_outputs = StructuredOutputsParams(json=sample_json_schema)
+
+    test_sampling_params = [
+        dict(),
+        dict(frequency_penalty=-1.0),
+        dict(bad_words=["the", " the"]),
+        dict(logprobs=2),
+        dict(logprobs=2, frequency_penalty=-1.0),
+        dict(structured_outputs=struct_outputs),
+        dict(
+            structured_outputs=struct_outputs,
+            logprobs=2,
+            frequency_penalty=-1.0,
+        ),
+    ]
+
+    # test_preemption, executor, async_scheduling,
+    # spec_config, test_prefill_chunking
+    test_configs = [
+        (False, "mp", False, None, False),
+        (False, "mp", False, spec_config, False),
+        (True, "mp", False, spec_config, True),
+        (True, "uni", False, spec_config_short, True),
+        (False, "mp", True, spec_config, False),
+        (True, "mp", True, spec_config, False),
+        (False, "mp", True, spec_config_short, True),
+        (True, "uni", True, spec_config, False),
+        (True, "uni", True, spec_config_short, False),
+        (True, "mp", True, spec_config, True),
+        (True, "uni", True, spec_config_short, True),
+    ]
+
+    run_tests(monkeypatch, MTP_MODEL, test_configs, test_sampling_params)
+
+
+@pytest.mark.flaky(reruns=2, only_on=current_platform.is_rocm())
+def test_with_ngram_gpu_spec_decoding(monkeypatch: pytest.MonkeyPatch):
+    """Test ngram_gpu speculative decoding with different configurations.
+
+    This test specifically validates ngram_gpu behavior with various:
+    - Number of speculative tokens (2-6)
+    - Prompt lookup window sizes (min/max)
+    - Async scheduling enabled (as in production)
+    - Different executors and chunking settings
+    """
+
+    # Variant with larger speculation window
+    ngram_gpu_config = {
+        "method": "ngram_gpu",
+        "num_speculative_tokens": 3,
+        "prompt_lookup_max": 3,
+        "prompt_lookup_min": 2,
+    }
+
+    # Test configurations covering various scenarios
+    # test_preemption, executor, async_scheduling,
+    # spec_config, test_prefill_chunking
+    test_configs = [
+        (False, "mp", False, None, False),
+        (False, "mp", False, ngram_gpu_config, False),
+        (True, "mp", False, ngram_gpu_config, True),
+        (False, "mp", True, ngram_gpu_config, False),
+        (True, "mp", True, ngram_gpu_config, False),
+        (True, "uni", True, ngram_gpu_config, False),
+        (True, "mp", True, ngram_gpu_config, True),
+    ]
+
+    # Use MODEL (Qwen) for ngram_gpu tests as it's lighter weight
+    # and ngram_gpu doesn't require a specific draft model
+    run_tests(monkeypatch, MODEL, test_configs, [{}])
+
+
+@dynamo_config.patch(cache_size_limit=16)
+def run_tests(
+    monkeypatch: pytest.MonkeyPatch,
+    model: str,
+    test_configs: list[tuple],
+    test_sampling_params: list[dict[str, Any]],
+):
+    """Test consistency of combos of async scheduling, preemption,
+    uni/multiproc executor with spec decoding."""
+
+    # Flex attention supports float32.
+    attention_config = {"backend": "FLEX_ATTENTION"}
+
+    with monkeypatch.context() as m:
+        # lock matmul precision to full FP32 (IEEE)
+        m.setenv("VLLM_FLOAT32_MATMUL_PRECISION", "highest")
+        outputs: list[tuple[str, list, list]] = []
+        for n, (
+            test_preemption,
+            executor,
+            async_scheduling,
+            spec_config,
+            test_prefill_chunking,
+        ) in enumerate(test_configs, 1):
+            test_str = f"{n}/{len(test_configs)}"
+            test_results = run_test(
+                model,
+                test_str,
+                test_sampling_params,
+                test_preemption,
+                executor,
+                async_scheduling,
+                spec_config,
+                test_prefill_chunking=test_prefill_chunking,
+                attention_config=attention_config,
+            )
+            outputs.append(test_results)
+
+    baseline_config, baseline_tests, _ = outputs[0]
+    _, _, baseline_acceptances = next(
+        (o for o in outputs if o[2] is not None), (None, None, None)
+    )
+
+    print(f"BASELINE: config=[{baseline_config}], accept_rates={baseline_acceptances}")
+
+    failure = None
+    for test_config, test_outputs, test_acceptance_rates in outputs[1:]:
+        for (base_outs, base_logprobs), base_acceptance_rate, (
+            test_outs,
+            test_logprobs,
+        ), test_acceptance_rate, params in zip(
+            baseline_tests,
+            baseline_acceptances or repeat(None),
+            test_outputs,
+            test_acceptance_rates or repeat(None),
+            test_sampling_params,
+        ):
+            reason = None
+            try:
+                check_outputs_equal(
+                    outputs_0_lst=base_outs,
+                    outputs_1_lst=test_outs,
+                    name_0=f"baseline=[{baseline_config}], params={params}",
+                    name_1=f"config=[{test_config}], params={params}",
+                )
+            except AssertionError as e:
+                reason = "outputs ", e
+
+            if reason is None:
+                try:
+                    assert _all_logprobs_match(base_logprobs, test_logprobs)
+                except AssertionError as e:
+                    reason = "logprobs", e
+
+            if reason is None:
+                try:
+                    if (
+                        base_acceptance_rate is not None
+                        and test_acceptance_rate is not None
+                    ):
+                        if "spec_mml=None" in test_config:
+                            # Preemption causes more variance in acceptance rates
+                            if (
+                                current_platform.is_rocm()
+                                and "preemption=True" in test_config
+                            ):
+                                tolerance = 0.10
+                            else:
+                                tolerance = 0.05
+                            assert (
+                                test_acceptance_rate > base_acceptance_rate
+                                or test_acceptance_rate
+                                == pytest.approx(base_acceptance_rate, rel=tolerance)
+                            )
+                        else:
+                            # Currently the reported acceptance rate is expected to be
+                            # lower when we sometimes skip drafting altogether.
+                            assert test_acceptance_rate > 0.1
+                except AssertionError as e:
+                    reason = "accept  ", e
+
+            if reason is None:
+                print(
+                    f"\033[32mPASSED\033[0m:           "
+                    f"config=[{test_config}], params={params}"
+                    f" accept_rate={test_acceptance_rate}"
+                )
+            else:
+                reason_str, _ = reason
+                print(
+                    f"\033[31mFAILED\033[0m({reason_str}): "
+                    f"config=[{test_config}], params={params}"
+                    f" accept_rate={test_acceptance_rate}"
+                )
+                if failure is None:
+                    _, failure = reason
+
+    if failure is not None:
+        raise failure
+
+
+def run_test(
+    model: str,
+    test_str: str,
+    sampling_param_tests: list[dict[str, Any]],
+    test_preemption: bool,
+    executor: str,
+    async_scheduling: bool,
+    spec_config: dict[str, Any] | None,
+    test_prefill_chunking: bool,
+    attention_config: dict[str, Any] | None = None,
+):
+    spec_decoding = spec_config is not None
+    cache_arg: dict[str, Any] = (
+        # Force preemptions
+        dict(num_gpu_blocks_override=32)
+        if test_preemption
+        else dict(gpu_memory_utilization=0.9)
+    )
+    spec_mml = (spec_config or {}).get("max_model_len")
+    spec_method = (spec_config or {}).get("method", "none")
+    test_config = (
+        f"executor={executor}, preemption={test_preemption}, "
+        f"async_sched={async_scheduling}, "
+        f"chunk_prefill={test_prefill_chunking}, "
+        f"spec_decoding={spec_decoding}, spec_method={spec_method}, spec_mml={spec_mml}"
+    )
+    print("-" * 80)
+    print(f"---- TESTING {test_str}: {test_config}")
+    print("-" * 80)
+
+    with VllmRunner(
+        model,
+        max_model_len=4096,
+        enable_chunked_prefill=test_prefill_chunking,
+        # Force prefill chunking
+        max_num_batched_tokens=48 if test_prefill_chunking else None,
+        enforce_eager=ENFORCE_EAGER,
+        async_scheduling=async_scheduling,
+        distributed_executor_backend=executor,
+        dtype="float32",
+        speculative_config=spec_config,
+        disable_log_stats=False,
+        attention_config=attention_config,
+        enable_prefix_caching=False if current_platform.is_rocm() else None,
+        **cache_arg,
+    ) as vllm_model:
+        results = []
+        acceptance_rates: list[float] | None = [] if spec_decoding else None
+        for override_params in sampling_param_tests:
+            metrics_before = vllm_model.llm.get_metrics()
+            print(f"----------- RUNNING PARAMS: {override_params}")
+            results.append(
+                vllm_model.generate(
+                    example_prompts,
+                    sampling_params=SamplingParams(**default_params, **override_params),
+                    return_logprobs=True,
+                )
+            )
+            metrics_after = vllm_model.llm.get_metrics()
+            if acceptance_rates is not None:
+                acceptance_rate = _get_acceptance_rate(metrics_before, metrics_after)
+                acceptance_rates.append(acceptance_rate)
+                print(f"ACCEPTANCE RATE {acceptance_rate}")
+
+            if test_preemption:
+                preemptions = _get_count(
+                    metrics_before, metrics_after, "vllm:num_preemptions"
+                )
+                assert preemptions > 0, "preemption test had no preemptions"
+
+    if len(results) > 1:
+        # First check that the different parameter configs
+        # actually result in different output.
+        for (other_test_outs, other_test_logprobs), params in zip(
+            results[1:], sampling_param_tests[1:]
+        ):
+            with pytest.raises(AssertionError):
+                check_outputs_equal(
+                    outputs_0_lst=results[0][0],
+                    outputs_1_lst=other_test_outs,
+                    name_0=f"baseline params={params}",
+                    name_1=f"other params={params}",
+                )
+                assert _all_logprobs_match(results[0][1], other_test_logprobs)
+
+    return test_config, results, acceptance_rates
+
+
+def _all_logprobs_match(req_a, req_b) -> bool:
+    return (
+        req_a == req_b
+        or len(req_a) == len(req_b)
+        and all(
+            len(seq_a) == len(seq_b)
+            and all(_logprobs_match(a, b) for a, b in zip(seq_a, seq_b))
+            for seq_a, seq_b in zip(req_a, req_b)
+        )
+    )
+
+
+def _logprobs_match(lps_a: dict[int, Logprob], lps_b: dict[int, Logprob]) -> bool:
+    rel_tol, abs_tol = 1e-3, 1e-6
+    return (
+        len(lps_a) == len(lps_b)
+        and lps_a.keys() == lps_b.keys()
+        and all(
+            a.decoded_token == b.decoded_token
+            and a.rank == b.rank
+            and a.logprob == pytest.approx(b.logprob, rel=rel_tol, abs=abs_tol)
+            for a, b in ((lps_a[x], lps_b[x]) for x in lps_a)
+        )
+    )
+
+
+def _get_acceptance_rate(before: list[Metric], after: list[Metric]) -> float:
+    draft = _get_count(before, after, "vllm:spec_decode_num_draft_tokens")
+    accept = _get_count(before, after, "vllm:spec_decode_num_accepted_tokens")
+    return accept / draft if draft > 0 else 0.0
+
+
+def _get_count(before: list[Metric], after: list[Metric], name: str) -> int:
+    before_val = next(m.value for m in before if m.name == name)
+    after_val = next(m.value for m in after if m.name == name)
+    return after_val - before_val

third_party/vllm/tests/v1/e2e/general/test_cascade_attention.py

@@ -0,0 +1,36 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+import pytest
+
+from vllm import LLM, SamplingParams
+
+from ....utils import create_new_process_for_each_test
+
+
+@create_new_process_for_each_test()
+@pytest.mark.parametrize("attn_backend", ["FLASH_ATTN", "FLASHINFER"])
+def test_cascade_attention(example_system_message, attn_backend):
+    prompt = "\n<User>: Implement fibonacci sequence in Python.\n<Claude>:"
+
+    if attn_backend == "FLASHINFER":
+        pytest.skip(
+            "This test is failing with FlashInfer backend and "
+            "needs investigation. See issue #25679."
+        )
+
+    llm = LLM(
+        model="Qwen/Qwen2-1.5B-Instruct", attention_config={"backend": attn_backend}
+    )
+    sampling_params = SamplingParams(temperature=0.0, max_tokens=100)
+
+    # No cascade attention.
+    single_prompt = [example_system_message + prompt]
+    responses = llm.generate(single_prompt, sampling_params)
+    ref_output = responses[0].outputs[0].text
+
+    # (Probably) Use cascade attention.
+    prompts = [example_system_message + prompt] * 64
+    responses = llm.generate(prompts, sampling_params)
+    for response in responses:
+        assert response.outputs[0].text == ref_output

third_party/vllm/tests/v1/e2e/general/test_context_length.py

@@ -0,0 +1,63 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""
+Tests for vLLM `vllm/v1/engine/processor.Processor._validate_model_input()`
+handling of maximum context length for decoder models.
+
+This test ensures:
+- A prompt that is one token shorter than the model's maximum context length
+  can be processed successfully when requesting one additional token.
+- A prompt that reaches the model's maximum context length throws a
+  `ValueError` when requesting at least one additional token.
+"""
+
+import pytest
+
+from tests.conftest import VllmRunner
+from tests.utils import create_new_process_for_each_test
+
+
+@create_new_process_for_each_test()
+@pytest.mark.parametrize("model, max_model_len", [("JackFram/llama-160m", 2048)])
+@pytest.mark.parametrize(
+    "prompt_len, max_tokens",
+    [
+        (2047, 1),  # prompt_len = max_model_len - 1 -> allowed
+        (2048, 1),  # prompt_len = max_model_len -> not allowed
+    ],
+)
+def test_decoder_max_context_length_validation(
+    model: str,
+    max_model_len: int,
+    vllm_runner: type[VllmRunner],
+    prompt_len: int,
+    max_tokens: int,
+) -> None:
+    """Check vLLM decoder model input validation for edge cases where
+    the prompt length is (almost) equal to the max model length."""
+
+    prompt_ids = [[43] * prompt_len]
+
+    with vllm_runner(
+        model_name=model,
+        tokenizer_name=model,
+        max_model_len=max_model_len,
+        max_num_seqs=1,
+        tensor_parallel_size=1,
+    ) as vllm_model:
+        if prompt_len + max_tokens <= max_model_len:
+            # Should succeed as constraints are met
+            vllm_model.generate_greedy(prompt_ids, max_tokens)
+        else:
+            # Should raise the ValueError defined in
+            # vllm/v1/engine/processor.Processor_validate_model_input()
+            expected_msg = (
+                f"The decoder prompt (length {prompt_len}) plus the number of "
+                f"requested output tokens (at least 1) is longer than "
+                f"the maximum model length of {max_model_len}. "
+                "Make sure that `max_model_len` is no smaller than the number of "
+                "text tokens (prompt + requested output tokens)."
+            )
+            with pytest.raises(ValueError) as excinfo:
+                vllm_model.generate_greedy(prompt_ids, max_tokens)
+            assert expected_msg in str(excinfo.value)

third_party/vllm/tests/v1/e2e/general/test_correctness_sliding_window.py

@@ -0,0 +1,98 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+from dataclasses import dataclass
+
+import pytest
+
+from vllm import LLM, SamplingParams
+from vllm.platforms import current_platform
+
+from ....utils import check_answers, prep_prompts
+
+
+@dataclass
+class TestConfig:
+    sliding_window: int
+    ln_range: tuple[int, int]
+
+
+model_config = {
+    "bigcode/starcoder2-3b": TestConfig(4096, (800, 1100)),
+    "google/gemma-3-1b-it": TestConfig(4096, (400, 800)),
+}
+
+
+@pytest.mark.parametrize(
+    "model",
+    [
+        "bigcode/starcoder2-3b",  # sliding window only
+        "google/gemma-3-1b-it",  # sliding window + full attention
+    ],
+)
+@pytest.mark.parametrize("batch_size", [5])
+@pytest.mark.parametrize("seed", [1])
+@pytest.mark.parametrize("disable_hybrid_kv_cache_manager", [True, False])
+def test_sliding_window_retrieval(
+    model, batch_size, seed, disable_hybrid_kv_cache_manager
+):
+    """
+    The test does a bunch of assignments "x1 = 10\nx2 = 33\n..." and then
+    asks for value of one of them (which is outside the sliding window).
+    If we tell it upfront which we are going to be looking for, then
+    it answers correctly (mostly).
+    """
+    # NOTE: For ROCm, we have to enforce eager mode to use custom kernel
+    # implementation of GELU with tanh approximation, as PyTorch's native
+    # implementation is currently unstable with torch.compile and produces garbage.
+    enforce_eager = current_platform.is_rocm()
+
+    test_config = model_config[model]
+
+    llm = LLM(
+        model=model,
+        disable_hybrid_kv_cache_manager=disable_hybrid_kv_cache_manager,
+        enforce_eager=enforce_eager,
+    )
+    sampling_params = SamplingParams(temperature=0.0, max_tokens=100)
+
+    prompts, answer, indices = prep_prompts(batch_size, ln_range=test_config.ln_range)
+
+    check_length(prompts, llm, test_config.sliding_window)
+
+    # Fresh generation
+    responses = llm.generate(prompts, sampling_params)
+    check_answers(
+        indices,
+        answer,
+        [response.outputs[0].text for response in responses],
+        accept_rate=1.0,
+    )
+
+    # Re-generate with the same prompts to test prefix caching
+    responses = llm.generate(prompts, sampling_params)
+    check_answers(
+        indices,
+        answer,
+        [response.outputs[0].text for response in responses],
+        accept_rate=1.0,
+    )
+
+
+def check_length(prompts: list[str], llm: LLM, sliding_window: int):
+    """
+    Check if the prompt length is valid, i.e., longer than the sliding window
+    size and shorter than the model's max length.
+
+    Args:
+        prompts: list of prompts
+        llm: LLM object
+        sliding_window: Sliding window size
+    """
+    tokenizer = llm.get_tokenizer()
+    max_model_len = llm.llm_engine.model_config.max_model_len
+    assert any(len(tokenizer.encode(prompt)) > sliding_window for prompt in prompts), (
+        "Prompt is too short for test"
+    )
+    assert all(len(tokenizer.encode(prompt)) <= max_model_len for prompt in prompts), (
+        "Prompt is too long for test"
+    )

third_party/vllm/tests/v1/e2e/general/test_kv_sharing_fast_prefill.py

@@ -0,0 +1,102 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+import random
+
+import pytest
+
+from vllm import LLM, SamplingParams
+from vllm.config import CompilationConfig, CompilationMode
+from vllm.platforms import current_platform
+
+from ....utils import check_answers, fork_new_process_for_each_test, prep_prompts
+
+# global seed
+SEED = 42
+
+
+@pytest.fixture
+def test_prompts():
+    """
+    Adapted from tests/v1/e2e/spec_decode/test_spec_decode.py
+    """
+    prompt_types = ["repeat", "sentence"]
+    # Setting higher num prompts increases the chance of numerics mismatch
+    # due to matrix multiplication numerics depending on batch dimension
+    num_prompts = 10
+    prompts = []
+
+    random.seed(0)
+    random_prompt_type_choices = random.choices(prompt_types, k=num_prompts)
+
+    for kind in random_prompt_type_choices:
+        word_choices = ["test", "temp", "hello", "where"]
+        word = random.choice(word_choices)
+        if kind == "repeat":
+            prompt = f"""please repeat the word '{word}' 10 times."""
+        elif kind == "sentence":
+            prompt = f"""please give a ten-word sentence that
+            uses the word {word} at least once."""
+        else:
+            raise ValueError(f"Unknown prompt type: {kind}")
+        prompts.append(prompt)
+
+    return prompts
+
+
+use_fork_for_test = (
+    fork_new_process_for_each_test if not current_platform.is_rocm() else lambda x: x
+)
+
+
+@use_fork_for_test
+@pytest.mark.parametrize("kv_sharing_fast_prefill", [False, True])
+@pytest.mark.parametrize("enforce_eager", [True, False])
+def test_kv_sharing_fast_prefill(
+    monkeypatch: pytest.MonkeyPatch,
+    kv_sharing_fast_prefill: bool,
+    enforce_eager: bool,
+):
+    if not enforce_eager and current_platform.is_rocm():
+        # Relevant context: https://github.com/vllm-project/vllm/pull/29244
+        pytest.skip(
+            "ROCm: torch.compile produces incorrect output for gemma-3n's GELU "
+            "with tanh approximation. Use enforce_eager=True instead."
+        )
+
+    sampling_params = SamplingParams(temperature=0.0, max_tokens=100)
+    compilation_config = CompilationConfig(
+        # This allows vLLM compilation backend to handle allocating and
+        # managing buffers for cudagraph
+        cudagraph_copy_inputs=True,
+        mode=CompilationMode.VLLM_COMPILE
+        if not enforce_eager
+        else CompilationMode.NONE,
+    )
+    batch_size = 10
+
+    with monkeypatch.context() as m:
+        # Make scheduling deterministic for reproducibility
+        if current_platform.is_rocm():
+            # Use spawn to prevent cuda re-initialization error
+            m.setenv("VLLM_WORKER_MULTIPROC_METHOD", "spawn")
+        else:
+            m.setenv("VLLM_ENABLE_V1_MULTIPROCESSING", "0")
+
+        prompts, answer, indices = prep_prompts(batch_size)
+
+        llm = LLM(
+            model="google/gemma-3n-E2B-it",
+            enforce_eager=enforce_eager,
+            compilation_config=compilation_config,
+            seed=SEED,
+            kv_sharing_fast_prefill=kv_sharing_fast_prefill,
+            attention_backend="TRITON_ATTN",
+        )
+        responses = llm.generate(prompts, sampling_params)
+        check_answers(
+            indices,
+            answer,
+            [response.outputs[0].text for response in responses],
+            accept_rate=1.0,
+        )

third_party/vllm/tests/v1/e2e/general/test_mamba_prefix_cache.py

@@ -0,0 +1,809 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+import multiprocessing as mp
+import os
+import traceback
+from collections.abc import Callable
+from dataclasses import dataclass
+from typing import Any
+
+import datasets
+import pytest
+import torch
+
+from tests.utils import create_new_process_for_each_test
+from vllm import LLM, SamplingParams, TokensPrompt
+from vllm.config import CacheConfig
+from vllm.distributed import cleanup_dist_env_and_memory
+from vllm.model_executor.layers.mamba.mamba_utils import MambaStateCopyFunc
+from vllm.sequence import IntermediateTensors
+from vllm.v1.attention.backends.utils import CommonAttentionMetadata
+from vllm.v1.core.kv_cache_manager import KVCacheBlocks, KVCacheManager
+from vllm.v1.core.sched.output import SchedulerOutput
+from vllm.v1.engine.core_client import InprocClient
+from vllm.v1.kv_cache_interface import KVCacheConfig
+from vllm.v1.outputs import SamplerOutput
+from vllm.v1.request import Request
+from vllm.v1.sample.metadata import SamplingMetadata
+from vllm.v1.spec_decode.metadata import SpecDecodeMetadata
+from vllm.v1.worker import mamba_utils
+from vllm.v1.worker.gpu_input_batch import CachedRequestState
+from vllm.v1.worker.gpu_model_runner import GPUModelRunner
+from vllm.v1.worker.lora_model_runner_mixin import GPUInputBatch
+from vllm.v1.worker.mamba_utils import get_mamba_groups
+
+
+@dataclass
+class StepAction:
+    num_computed_tokens_start: int
+    num_scheduled_tokens: int
+    kv_cache_block_ids: list[int]  # [] to follow last step
+    preprocess_copy_idx: tuple[int, int]  # -1, -1 for no copy
+    postprocess_copy_idx: tuple[int, int]  # -1, -1 for no copy
+
+
+num_speculative_tokens = 3
+
+num_accepted_tokens = 1
+prompt_token_ids: list[int] = []
+MODEL = "Qwen/Qwen3-Next-80B-A3B-Instruct-FP8"
+BLOCK_SIZE = 560
+NUM_HIDDEN_LAYERS = 1
+cur_step_action_idx = 0
+cur_step_action: StepAction | None = None
+step_actions: list[StepAction] = []
+
+
+def get_fake_sample_fn() -> SamplerOutput:
+    def fake_sample_fn(
+        self: GPUModelRunner,
+        logits: torch.Tensor | None,
+        spec_decode_metadata: SpecDecodeMetadata | None,
+    ) -> SamplerOutput:
+        assert logits is not None
+        num_computed_tokens_cpu_tensor = self.input_batch.num_computed_tokens_cpu_tensor
+        num_computed_tokens = num_computed_tokens_cpu_tensor[0].item()
+        if num_computed_tokens < self.input_batch.num_prompt_tokens[0].item():
+            first_token_id_index = self.input_batch.num_prompt_tokens[0].item()
+        else:
+            first_token_id_index = num_computed_tokens + 1
+        if spec_decode_metadata is None:
+            return SamplerOutput(
+                sampled_token_ids=torch.tensor(
+                    [[prompt_token_ids[first_token_id_index]]],
+                    device="cuda",
+                    dtype=torch.int32,
+                ),
+                logprobs_tensors=None,
+            )
+        accepted_tokens = prompt_token_ids[
+            first_token_id_index : first_token_id_index
+            + min(num_accepted_tokens, logits.shape[0])
+        ]
+        sampled_token_ids = accepted_tokens
+        return SamplerOutput(
+            sampled_token_ids=torch.tensor(
+                [sampled_token_ids], device="cuda", dtype=torch.int32
+            ),
+            logprobs_tensors=None,
+        )
+
+    return fake_sample_fn
+
+
+def get_fake_propose_draft_token_ids_fn():
+    def fake_propose_draft_token_ids_fn(
+        self: GPUModelRunner,
+        scheduler_output: SchedulerOutput,
+        sampled_token_ids: torch.Tensor | list[list[int]],
+        sampling_metadata: SamplingMetadata,
+        hidden_states: torch.Tensor,
+        sample_hidden_states: torch.Tensor,
+        aux_hidden_states: list[torch.Tensor] | None,
+        spec_decode_metadata: SpecDecodeMetadata | None,
+        common_attn_metadata: CommonAttentionMetadata,
+        slot_mappings: dict[str, torch.Tensor] | list[dict[str, torch.Tensor]] | None,
+    ) -> list[list[int]]:
+        num_computed_tokens_cpu_tensor = self.input_batch.num_computed_tokens_cpu_tensor
+        num_computed_tokens = num_computed_tokens_cpu_tensor[0].item()
+        if (
+            self.input_batch.num_tokens_no_spec[0].item()
+            <= self.input_batch.num_prompt_tokens[0].item()
+        ):
+            first_token_id_index = self.input_batch.num_prompt_tokens[0].item()
+        else:
+            first_token_id_index = (
+                num_computed_tokens + 1
+            )  # bonus token isn't considered as computed
+        first_token_id_index += self.input_batch.num_accepted_tokens_cpu[0].item()
+        proposed_draft_token_ids = [
+            prompt_token_ids[
+                first_token_id_index : first_token_id_index + num_speculative_tokens
+            ]
+        ]
+
+        next_token_ids = torch.tensor(
+            prompt_token_ids[
+                first_token_id_index - 1 : first_token_id_index
+                - 1
+                + num_accepted_tokens
+            ],
+            device="cuda",
+            dtype=torch.int32,
+        )
+
+        valid_sampled_tokens_count = torch.tensor(
+            [num_accepted_tokens], device="cuda", dtype=torch.int32
+        )
+
+        self._copy_valid_sampled_token_count(next_token_ids, valid_sampled_tokens_count)
+
+        return torch.tensor(proposed_draft_token_ids, device="cuda", dtype=torch.int32)
+
+    return fake_propose_draft_token_ids_fn
+
+
+def get_fake_step_action_fn(original_step_action_fn: Callable):
+    def fake_get_output(self: InprocClient):
+        global cur_step_action_idx
+        global cur_step_action
+        if cur_step_action_idx < len(step_actions):
+            cur_step_action = step_actions[cur_step_action_idx]
+            cur_step_action_idx += 1
+        else:
+            cur_step_action = None
+        print(f"cur_step_action: {cur_step_action_idx=} {cur_step_action=}")
+        return original_step_action_fn(self)
+
+    return fake_get_output
+
+
+def get_fake_allocate_slots_fn(original_allocate_slots_fn: Callable):
+    def fake_allocate_slots_fn(
+        self: KVCacheManager,
+        request: Request,
+        num_new_tokens: int,
+        num_new_computed_tokens: int = 0,
+        new_computed_blocks: KVCacheBlocks | None = None,
+        num_lookahead_tokens: int = 0,
+        num_external_computed_tokens: int = 0,
+        delay_cache_blocks: bool = False,
+        num_encoder_tokens: int = 0,
+    ):
+        ret = original_allocate_slots_fn(
+            self,
+            request,
+            num_new_tokens,
+            num_new_computed_tokens,
+            new_computed_blocks,
+            num_lookahead_tokens,
+            num_external_computed_tokens,
+            delay_cache_blocks,
+            num_encoder_tokens,
+        )
+        if cur_step_action is not None:
+            cur_block_ids = self.coordinator.single_type_managers[0].req_to_blocks[
+                request.request_id
+            ]
+            not_null_block_flags = [not block.is_null for block in cur_block_ids]
+            block_ids = [1 if block else 0 for block in not_null_block_flags]
+            assert block_ids == cur_step_action.kv_cache_block_ids
+        return ret
+
+    return fake_allocate_slots_fn
+
+
+mamba_kv_cache_dict = {}
+
+
+def get_fake_execute_model_fn(original_execute_model_fn: Callable):
+    last_num_computed_tokens = 0
+    num_prompt_tokens = None
+
+    def fake_execute_model_fn(
+        self: GPUModelRunner,
+        scheduler_output: SchedulerOutput,
+        intermediate_tensors: IntermediateTensors | None = None,
+    ):
+        if cur_step_action is not None:
+            num_scheduled_tokens = next(
+                iter(scheduler_output.num_scheduled_tokens.values())
+            )
+            assert num_scheduled_tokens == cur_step_action.num_scheduled_tokens
+        mamba_group_ids, mamba_spec = get_mamba_groups(self.kv_cache_config)
+        mamba_group_id = mamba_group_ids[0]
+        mamba_layer_name = self.kv_cache_config.kv_cache_groups[
+            mamba_group_id
+        ].layer_names[0]
+        nonlocal last_num_computed_tokens
+        nonlocal num_prompt_tokens
+
+        if (
+            len(scheduler_output.scheduled_new_reqs) > 0
+            and scheduler_output.scheduled_new_reqs[0].prompt_token_ids is not None
+        ):
+            # record number of prompt tokens
+            num_prompt_tokens = len(
+                scheduler_output.scheduled_new_reqs[0].prompt_token_ids
+            )
+
+        if len(scheduler_output.scheduled_cached_reqs.req_ids) > 0:
+            num_computed_tokens = (
+                scheduler_output.scheduled_cached_reqs.num_computed_tokens[0]
+            )
+            if (
+                self.num_spec_tokens
+                and num_prompt_tokens is not None
+                and num_computed_tokens > num_prompt_tokens
+            ):
+                # NOTE (tdoublep) with async scheduling, the scheduler does not have an
+                # accurate measure of the number of computed tokens; we need to subtract
+                # the number of reject tokens from the previous timestep.
+                num_computed_tokens -= num_speculative_tokens + 1 - num_accepted_tokens
+            if (
+                num_computed_tokens // BLOCK_SIZE
+                > last_num_computed_tokens // BLOCK_SIZE
+            ):
+                # generated a new aligned block in this step
+                block_idx = num_computed_tokens // mamba_spec.block_size - 1
+                block_id = (
+                    self.input_batch.block_table.block_tables[mamba_group_id]
+                    .block_table.cpu[0, block_idx]
+                    .item()
+                )
+                if block_id != 0:
+                    kv_cache = self.compilation_config.static_forward_context[
+                        mamba_layer_name
+                    ].kv_cache
+                    mamba_kv_cache_dict[
+                        num_computed_tokens - num_computed_tokens % BLOCK_SIZE
+                    ] = (
+                        kv_cache[0][0][block_id].clone(),
+                        kv_cache[0][1][block_id].clone(),
+                    )
+
+            last_num_computed_tokens = num_computed_tokens
+        else:
+            last_num_computed_tokens = 0
+
+        ret = original_execute_model_fn(self, scheduler_output, intermediate_tensors)
+
+        if cur_step_action is not None:
+            assert (
+                cur_step_action.num_computed_tokens_start
+                == self.input_batch.num_computed_tokens_cpu[0].item()
+            )
+
+        return ret
+
+    return fake_execute_model_fn
+
+
+def get_fake_process_mamba_fn(
+    original_preprocess_mamba_fn: Callable,
+    original_post_process_mamba_fn: Callable,
+    original_copy_fn: Callable,
+):
+    copy_info: tuple[list[int], list[int], list[int]] | None = None
+
+    def check_copy_info(
+        action: tuple[int, int],
+        kv_cache_config: KVCacheConfig,
+        forward_context: dict[str, Any],
+        input_batch: GPUInputBatch,
+    ):
+        assert copy_info is not None
+        if action == (-1, -1):
+            assert len(copy_info[0]) == len(copy_info[1]) == len(copy_info[2]) == 0
+        else:
+            assert len(copy_info[0]) == len(copy_info[1]) == len(copy_info[2]) == 2
+            mamba_group_ids, mamba_spec = get_mamba_groups(kv_cache_config)
+            mamba_group_id = mamba_group_ids[0]
+            mamba_layer_name = kv_cache_config.kv_cache_groups[
+                mamba_group_id
+            ].layer_names[0]
+            mamba_kv_cache = forward_context[mamba_layer_name].kv_cache[0][-1]
+            mamba_block_table = input_batch.block_table.block_tables[
+                mamba_group_id
+            ].block_table.cpu[0]
+            expected_temporal_src = mamba_kv_cache[
+                mamba_block_table[action[0]]
+            ].data_ptr()
+            expected_temporal_dest = mamba_kv_cache[
+                mamba_block_table[action[1]]
+            ].data_ptr()
+            # -1 is qwen3-next's temporal. We skip checking conv as it is more complex.
+            assert copy_info[0][-1] == expected_temporal_src
+            assert copy_info[1][-1] == expected_temporal_dest
+
+    def fake_preprocess_mamba_fn(
+        scheduler_output: SchedulerOutput,
+        kv_cache_config: KVCacheConfig,
+        cache_config: CacheConfig,
+        mamba_state_idx: dict[str, int],
+        input_batch: GPUInputBatch,
+        requests: dict[str, CachedRequestState],
+        forward_context: dict[str, Any],
+        mamba_state_copy_funcs: tuple[MambaStateCopyFunc, ...],
+        copy_bufs: mamba_utils.MambaCopyBuffers,
+    ):
+        nonlocal copy_info
+        copy_info = None
+        ret = original_preprocess_mamba_fn(
+            scheduler_output,
+            kv_cache_config,
+            cache_config,
+            mamba_state_idx,
+            input_batch,
+            requests,
+            forward_context,
+            mamba_state_copy_funcs,
+            copy_bufs,
+        )
+        if cur_step_action is not None:
+            check_copy_info(
+                cur_step_action.preprocess_copy_idx,
+                kv_cache_config,
+                forward_context,
+                input_batch,
+            )
+        return ret
+
+    def fake_post_process_mamba_fn(
+        scheduler_output: SchedulerOutput,
+        kv_cache_config: KVCacheConfig,
+        input_batch: GPUInputBatch,
+        requests: dict[str, CachedRequestState],
+        mamba_state_idx: dict[str, int],
+        forward_context: dict[str, Any],
+        mamba_state_copy_funcs: tuple[MambaStateCopyFunc, ...],
+        copy_bufs: mamba_utils.MambaCopyBuffers,
+    ):
+        nonlocal copy_info
+        copy_info = None
+        ret = original_post_process_mamba_fn(
+            scheduler_output,
+            kv_cache_config,
+            input_batch,
+            requests,
+            mamba_state_idx,
+            forward_context,
+            mamba_state_copy_funcs,
+            copy_bufs,
+        )
+        if cur_step_action is not None:
+            check_copy_info(
+                cur_step_action.postprocess_copy_idx,
+                kv_cache_config,
+                forward_context,
+                input_batch,
+            )
+        return ret
+
+    def fake_copy_fn(copy_bufs: mamba_utils.MambaCopyBuffers):
+        nonlocal copy_info
+        assert copy_info is None
+        n = copy_bufs.offset
+        src_state_list = copy_bufs.src_ptrs.cpu[:n].tolist()
+        dest_state_list = copy_bufs.dst_ptrs.cpu[:n].tolist()
+        num_elements_list = copy_bufs.sizes.cpu[:n].tolist()
+        copy_info = (src_state_list, dest_state_list, num_elements_list)
+        return original_copy_fn(copy_bufs)
+
+    return fake_preprocess_mamba_fn, fake_post_process_mamba_fn, fake_copy_fn
+
+
+def run_ref_mamba_state_in_subprocess() -> None:
+    ctx = mp.get_context("spawn")
+    proc = ctx.Process(target=_run_ref_mamba_state_worker)
+    proc.start()
+    proc.join(timeout=600)
+    if proc.exitcode != 0:
+        raise RuntimeError(f"Ref mamba state process exited with code {proc.exitcode}.")
+
+
+def _run_ref_mamba_state_worker():
+    try:
+        os.environ["VLLM_ENABLE_V1_MULTIPROCESSING"] = "0"
+        num_generated_tokens = 8000
+        num_prompt_tokens = 500
+        sampling_params = SamplingParams(
+            temperature=0.0, max_tokens=num_generated_tokens
+        )
+        prompt_dataset = datasets.load_dataset("heheda/a_long_article")
+        full_prompt = prompt_dataset["train"][0]["text"]
+        fake_execute_model_fn = get_fake_execute_model_fn(GPUModelRunner.execute_model)
+        GPUModelRunner.execute_model = fake_execute_model_fn
+        fake_sample_fn = get_fake_sample_fn()
+        GPUModelRunner._sample = fake_sample_fn
+        engine = LLM(
+            model=MODEL,
+            block_size=BLOCK_SIZE,
+            hf_overrides={"num_hidden_layers": NUM_HIDDEN_LAYERS},
+            seed=42,
+        )
+        global prompt_token_ids
+        prompt_token_ids = engine.get_tokenizer().encode(full_prompt)
+        print(f"Token IDs length: {len(prompt_token_ids)}")
+
+        _outputs = engine.generate(
+            [TokensPrompt(prompt_token_ids=prompt_token_ids[:num_prompt_tokens])],
+            sampling_params,
+        )
+        # ref_mamba_kv_cache_dict = torch.load("mamba_kv_cache_dict.pth")
+        # check_mamba_state_equal(ref_mamba_kv_cache_dict, mamba_kv_cache_dict)
+        # torch.save(mamba_kv_cache_dict, "mamba_kv_cache_dict.pth")
+        cpu_state_ref = {
+            key: tuple(tensor.detach().cpu() for tensor in tensors)
+            for key, tensors in mamba_kv_cache_dict.items()
+        }
+        torch.save(cpu_state_ref, "mamba_kv_cache_dict_ref.pth")
+        mamba_kv_cache_dict.clear()
+        del engine
+        torch.accelerator.empty_cache()
+        cleanup_dist_env_and_memory()
+    except Exception:
+        traceback.print_exc()
+        raise
+
+
+def check_mamba_state_equal(
+    mamba_state_ref: dict, mamba_state_new: dict, keys_to_check: list[int]
+):
+    atol = 1e-2
+    rtol = 1e-2
+    for key in keys_to_check:
+        assert key in mamba_state_new
+        assert key in mamba_state_ref
+        # mamba state new is a subset of mamba state ref
+        for i, (ref, new) in enumerate(zip(mamba_state_ref[key], mamba_state_new[key])):
+            if ref.device != new.device:
+                new = new.to(ref.device)
+            new = new[: ref.shape[0]]
+            if not torch.allclose(ref, new, atol=atol, rtol=rtol):
+                diff_mask = ~torch.isclose(ref, new, atol=atol, rtol=rtol)
+                diff_idx = torch.nonzero(diff_mask)
+                if diff_idx.shape[0] * 100 < ref.numel():
+                    print(
+                        f"[WARNING] found {diff_idx.shape[0] * 100 / ref.numel()}% of the elements are different"  # noqa: E501
+                    )
+                    continue
+                raise ValueError(
+                    f"Mamba state is not equal for key: {key} at index {i}"
+                )
+    return True
+
+
+@dataclass
+class TestConfig:
+    num_prompt_tokens: int
+    num_generated_tokens: int
+    num_accepted_tokens: int
+    step_actions: list[StepAction]
+
+
+def apply_patch(monkeypatch: pytest.MonkeyPatch):
+    monkeypatch.setenv("VLLM_ENABLE_V1_MULTIPROCESSING", "0")
+
+    fake_sample_fn = get_fake_sample_fn()
+    monkeypatch.setattr(GPUModelRunner, "_sample", fake_sample_fn)
+
+    fake_propose_draft_token_ids_fn = get_fake_propose_draft_token_ids_fn()
+    monkeypatch.setattr(
+        GPUModelRunner, "propose_draft_token_ids", fake_propose_draft_token_ids_fn
+    )
+
+    fake_execute_model_fn = get_fake_execute_model_fn(GPUModelRunner.execute_model)
+    monkeypatch.setattr(GPUModelRunner, "execute_model", fake_execute_model_fn)
+
+    fake_step_action_fn = get_fake_step_action_fn(InprocClient.get_output)
+    monkeypatch.setattr(InprocClient, "get_output", fake_step_action_fn)
+
+    fake_allocate_slots_fn = get_fake_allocate_slots_fn(KVCacheManager.allocate_slots)
+    monkeypatch.setattr(KVCacheManager, "allocate_slots", fake_allocate_slots_fn)
+
+    fake_preprocess_mamba_fn, fake_post_process_mamba_fn, fake_copy_fn = (
+        get_fake_process_mamba_fn(
+            mamba_utils.preprocess_mamba,
+            mamba_utils.postprocess_mamba,
+            mamba_utils.do_mamba_copy_block,
+        )
+    )
+    monkeypatch.setattr(mamba_utils, "preprocess_mamba", fake_preprocess_mamba_fn)
+    monkeypatch.setattr(mamba_utils, "postprocess_mamba", fake_post_process_mamba_fn)
+    monkeypatch.setattr(mamba_utils, "do_mamba_copy_block", fake_copy_fn)
+
+
+@create_new_process_for_each_test()
+def test_mamba_prefix_cache(monkeypatch: pytest.MonkeyPatch):
+    run_ref_mamba_state_in_subprocess()
+    apply_patch(monkeypatch)
+    prompt_dataset = datasets.load_dataset("heheda/a_long_article")
+    full_prompt = prompt_dataset["train"][0]["text"]
+    tests = {
+        "accept_1": TestConfig(
+            num_prompt_tokens=554,
+            num_generated_tokens=20,
+            num_accepted_tokens=1,
+            step_actions=[
+                StepAction(0, 554, [1, 1, 1, 1], (-1, -1), (-1, -1)),
+                StepAction(554, 4, [], (-1, -1), (-1, -1)),
+                StepAction(555, 4, [1, 1, 1, 1, 1], (-1, -1), (-1, -1)),
+                StepAction(556, 4, [], (-1, -1), (-1, -1)),
+                StepAction(557, 4, [], (0, 1), (-1, -1)),
+                StepAction(558, 4, [], (-1, -1), (-1, -1)),
+                StepAction(559, 4, [], (-1, -1), (1, 0)),
+                StepAction(560, 4, [], (-1, -1), (-1, -1)),
+                StepAction(561, 4, [0, 1, 1, 1, 1], (-1, -1), (-1, -1)),
+            ],
+        ),
+        # test case 2.1: no hit, accept 2 tokens
+        "accept_2_1": TestConfig(
+            num_prompt_tokens=554,
+            num_generated_tokens=20,
+            num_accepted_tokens=2,
+            step_actions=[
+                StepAction(0, 554, [1, 1, 1, 1], (-1, -1), (-1, -1)),
+                StepAction(554, 4, [], (-1, -1), (-1, -1)),
+                StepAction(556, 4, [1, 1, 1, 1, 1], (-1, -1), (-1, -1)),
+                StepAction(558, 4, [], (1, 1), (2, 0)),
+                StepAction(560, 4, [], (-1, -1), (-1, -1)),
+                StepAction(562, 4, [0, 1, 1, 1, 1], (-1, -1), (-1, -1)),
+            ],
+        ),
+        # test case 2.2: no hit, accept 2 tokens
+        "accept_2_2": TestConfig(
+            num_prompt_tokens=555,
+            num_generated_tokens=20,
+            num_accepted_tokens=2,
+            step_actions=[
+                StepAction(0, 555, [1, 1, 1, 1], (-1, -1), (-1, -1)),
+                StepAction(555, 4, [], (-1, -1), (-1, -1)),
+                StepAction(557, 4, [1, 1, 1, 1, 1], (1, 1), (-1, -1)),
+                StepAction(559, 4, [], (-1, -1), (1, 0)),
+                StepAction(561, 4, [], (-1, -1), (-1, -1)),
+                StepAction(563, 4, [0, 1, 1, 1, 1], (-1, -1), (-1, -1)),
+            ],
+        ),
+        "accept_3_1": TestConfig(
+            num_prompt_tokens=553,
+            num_generated_tokens=20,
+            num_accepted_tokens=3,
+            step_actions=[
+                StepAction(0, 553, [1, 1, 1, 1], (-1, -1), (-1, -1)),
+                StepAction(553, 4, [], (-1, -1), (-1, -1)),
+                StepAction(556, 4, [1, 1, 1, 1, 1], (-1, -1), (-1, -1)),
+                StepAction(559, 4, [], (2, 1), (1, 0)),
+                StepAction(562, 4, [], (-1, -1), (-1, -1)),
+                StepAction(565, 4, [0, 1, 1, 1, 1], (-1, -1), (-1, -1)),
+            ],
+        ),
+        "accept_3_2": TestConfig(
+            num_prompt_tokens=554,
+            num_generated_tokens=20,
+            num_accepted_tokens=3,
+            step_actions=[
+                StepAction(0, 554, [1, 1, 1, 1], (-1, -1), (-1, -1)),
+                StepAction(554, 4, [], (-1, -1), (-1, -1)),
+                StepAction(557, 4, [1, 1, 1, 1, 1], (2, 1), (3, 0)),
+                StepAction(560, 4, [], (-1, -1), (-1, -1)),
+                StepAction(563, 4, [0, 1, 1, 1, 1], (-1, -1), (-1, -1)),
+            ],
+        ),
+        "accept_3_3": TestConfig(
+            num_prompt_tokens=555,
+            num_generated_tokens=20,
+            num_accepted_tokens=3,
+            step_actions=[
+                StepAction(0, 555, [1, 1, 1, 1], (-1, -1), (-1, -1)),
+                StepAction(555, 4, [], (-1, -1), (-1, -1)),
+                StepAction(558, 4, [1, 1, 1, 1, 1], (2, 1), (2, 0)),
+                StepAction(561, 4, [], (-1, -1), (-1, -1)),
+                StepAction(564, 4, [0, 1, 1, 1, 1], (-1, -1), (-1, -1)),
+            ],
+        ),
+        "accept_4_1": TestConfig(
+            num_prompt_tokens=553,
+            num_generated_tokens=20,
+            num_accepted_tokens=4,
+            step_actions=[
+                StepAction(0, 553, [1, 1, 1, 1], (-1, -1), (-1, -1)),
+                StepAction(553, 4, [], (-1, -1), (-1, -1)),
+                StepAction(557, 4, [1, 1, 1, 1, 1], (3, 1), (3, 0)),
+                StepAction(561, 4, [], (-1, -1), (-1, -1)),
+                StepAction(565, 4, [0, 1, 1, 1, 1], (-1, -1), (-1, -1)),
+            ],
+        ),
+        "accept_4_2": TestConfig(
+            num_prompt_tokens=554,
+            num_generated_tokens=25,
+            num_accepted_tokens=4,
+            step_actions=[
+                StepAction(0, 554, [1, 1, 1, 1], (-1, -1), (-1, -1)),
+                StepAction(554, 4, [], (-1, -1), (-1, -1)),
+                StepAction(558, 4, [1, 1, 1, 1, 1], (3, 1), (2, 0)),
+                StepAction(562, 4, [], (-1, -1), (-1, -1)),
+                StepAction(566, 4, [0, 1, 1, 1, 1], (-1, -1), (-1, -1)),
+            ],
+        ),
+        "accept_4_3": TestConfig(
+            num_prompt_tokens=555,
+            num_generated_tokens=25,
+            num_accepted_tokens=4,
+            step_actions=[
+                StepAction(0, 555, [1, 1, 1, 1], (-1, -1), (-1, -1)),
+                StepAction(555, 4, [], (-1, -1), (-1, -1)),
+                StepAction(559, 4, [1, 1, 1, 1, 1], (3, 1), (1, 0)),
+                StepAction(563, 4, [], (-1, -1), (-1, -1)),
+                StepAction(567, 4, [0, 1, 1, 1, 1], (-1, -1), (-1, -1)),
+            ],
+        ),
+        "accept_4_4": TestConfig(
+            num_prompt_tokens=556,
+            num_generated_tokens=25,
+            num_accepted_tokens=4,
+            step_actions=[
+                StepAction(0, 556, [1, 1, 1, 1], (-1, -1), (-1, -1)),
+                StepAction(556, 4, [], (-1, -1), (3, 0)),
+                StepAction(560, 4, [1, 1, 1, 1, 1], (0, 1), (-1, -1)),
+                StepAction(564, 4, [0, 1, 1, 1, 1], (-1, -1), (-1, -1)),
+            ],
+        ),
+        "prompt_block_size": TestConfig(
+            num_prompt_tokens=560,
+            num_generated_tokens=10,
+            num_accepted_tokens=4,
+            step_actions=[
+                StepAction(0, 560, [1, 1, 1, 1], (-1, -1), (-1, -1)),
+                StepAction(560, 4, [1, 1, 1, 1, 1], (0, 1), (-1, -1)),
+            ],
+        ),
+        "prompt_2_block_size": TestConfig(
+            num_prompt_tokens=560 * 2,
+            num_generated_tokens=10,
+            num_accepted_tokens=4,
+            step_actions=[
+                StepAction(0, 560, [1, 1, 1, 1], (-1, -1), (-1, -1)),
+                StepAction(560, 560, [1, 1, 1, 1, 1], (0, 1), (-1, -1)),
+                StepAction(560 * 2, 4, [0, 1, 1, 1, 1, 1], (1, 2), (-1, -1)),
+            ],
+        ),
+        "prompt_2_block_size_10": TestConfig(
+            num_prompt_tokens=560 * 2 + 10,
+            num_generated_tokens=10,
+            num_accepted_tokens=4,
+            step_actions=[
+                StepAction(0, 560, [1, 1, 1, 1], (-1, -1), (-1, -1)),
+                StepAction(560, 570, [1, 0, 1, 1, 1, 1], (0, 2), (-1, -1)),
+                StepAction(560 * 2 + 10, 4, [0, 0, 1, 1, 1, 1], (-1, -1), (-1, -1)),
+            ],
+        ),
+        "prompt_3_block_size": TestConfig(
+            num_prompt_tokens=560 * 3,
+            num_generated_tokens=10,
+            num_accepted_tokens=4,
+            step_actions=[
+                StepAction(0, 560 * 2, [0, 1, 1, 1, 1], (-1, -1), (-1, -1)),
+                StepAction(560 * 2, 560, [0, 1, 1, 1, 1, 1], (1, 2), (-1, -1)),
+                StepAction(560 * 3, 4, [0, 0, 1, 1, 1, 1, 1], (2, 3), (-1, -1)),
+            ],
+        ),
+        "prompt_3_block_size_10": TestConfig(
+            num_prompt_tokens=560 * 3 + 10,
+            num_generated_tokens=10,
+            num_accepted_tokens=4,
+            step_actions=[
+                StepAction(0, 560 * 2, [0, 1, 1, 1, 1], (-1, -1), (-1, -1)),
+                StepAction(560 * 2, 570, [0, 1, 0, 1, 1, 1, 1], (1, 3), (-1, -1)),
+                StepAction(560 * 3 + 10, 4, [0, 0, 0, 1, 1, 1, 1], (-1, -1), (-1, -1)),
+            ],
+        ),
+        "prompt_10_block_size": TestConfig(
+            num_prompt_tokens=560 * 10,
+            num_generated_tokens=10,
+            num_accepted_tokens=4,
+            step_actions=[
+                StepAction(0, 560 * 5, [0, 0, 0, 0, 1, 1, 1, 1], (-1, -1), (-1, -1)),
+                StepAction(
+                    560 * 5,
+                    560 * 4,
+                    [0, 0, 0, 0, 1, 0, 0, 0, 1, 1, 1, 1],
+                    (4, 8),
+                    (-1, -1),
+                ),
+                StepAction(
+                    560 * 9,
+                    560,
+                    [0, 0, 0, 0, 0, 0, 0, 0, 1, 1, 1, 1, 1],
+                    (8, 9),
+                    (-1, -1),
+                ),
+                StepAction(
+                    560 * 10,
+                    4,
+                    [0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 1, 1, 1, 1],
+                    (9, 10),
+                    (-1, -1),
+                ),
+            ],
+        ),
+        "prompt_10_block_size_10": TestConfig(
+            num_prompt_tokens=560 * 10 + 10,
+            num_generated_tokens=10,
+            num_accepted_tokens=4,
+            step_actions=[
+                StepAction(0, 560 * 5, [0, 0, 0, 0, 1, 1, 1, 1], (-1, -1), (-1, -1)),
+                StepAction(
+                    560 * 5,
+                    560 * 4,
+                    [0, 0, 0, 0, 1, 0, 0, 0, 1, 1, 1, 1],
+                    (4, 8),
+                    (-1, -1),
+                ),
+                StepAction(
+                    560 * 9,
+                    560 + 10,
+                    [0, 0, 0, 0, 0, 0, 0, 0, 1, 0, 1, 1, 1, 1],
+                    (8, 10),
+                    (-1, -1),
+                ),
+            ],
+        ),
+    }
+
+    engine = LLM(
+        model=MODEL,
+        enable_prefix_caching=True,
+        block_size=BLOCK_SIZE,
+        mamba_cache_mode="align",
+        speculative_config={
+            "method": "qwen3_next_mtp",
+            "num_speculative_tokens": num_speculative_tokens,
+        },
+        max_num_batched_tokens=3072,
+        hf_overrides={"num_hidden_layers": NUM_HIDDEN_LAYERS},
+        seed=42,
+    )
+    global prompt_token_ids
+    prompt_token_ids = engine.get_tokenizer().encode(full_prompt)
+    print(f"Token IDs length: {len(prompt_token_ids)}")
+    for test_case_name, test_config in tests.items():
+        print(f"Running test case: {test_case_name}")
+        num_generated_tokens = test_config.num_generated_tokens
+        num_prompt_tokens = test_config.num_prompt_tokens
+        global num_accepted_tokens
+        num_accepted_tokens = test_config.num_accepted_tokens
+        sampling_params = SamplingParams(
+            temperature=0.0, max_tokens=num_generated_tokens
+        )
+        global cur_step_action_idx
+        cur_step_action_idx = 0
+        for step_action_prev, step_action_next in zip(
+            test_config.step_actions[:-1], test_config.step_actions[1:]
+        ):
+            if (
+                step_action_next.kv_cache_block_ids is not None
+                and len(step_action_next.kv_cache_block_ids) == 0
+            ):
+                prev_block_ids = step_action_prev.kv_cache_block_ids
+                if prev_block_ids is not None:
+                    step_action_next.kv_cache_block_ids = prev_block_ids.copy()
+        global step_actions
+        step_actions = test_config.step_actions
+        _ = engine.generate(
+            [TokensPrompt(prompt_token_ids=prompt_token_ids[:num_prompt_tokens])],
+            sampling_params,
+        )
+        assert engine.llm_engine.engine_core.engine_core.scheduler.reset_prefix_cache()
+        print(f"End test case: {test_case_name}")
+        keys_to_check = [
+            (action.postprocess_copy_idx[1] + 1) * BLOCK_SIZE
+            for action in test_config.step_actions
+            if action.postprocess_copy_idx and action.postprocess_copy_idx[0] != -1
+        ]
+        mamba_state_ref = torch.load("mamba_kv_cache_dict_ref.pth")
+        check_mamba_state_equal(mamba_state_ref, mamba_kv_cache_dict, keys_to_check)
+        mamba_kv_cache_dict.clear()
+    del engine
+    torch.accelerator.empty_cache()
+    cleanup_dist_env_and_memory()

third_party/vllm/tests/v1/e2e/general/test_min_tokens.py

@@ -0,0 +1,502 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""
+Comprehensive end-to-end tests for `min_tokens` in the V1 engine.
+
+Addresses #21950: verify and add CI coverage.
+
+Covers:
+1) Basic functionality
+2) Stop strings with `min_tokens` (bug #21987; fix in PR #22014)
+3) EOS behavior with `min_tokens` (potential logits-processor bug)
+4) Edge cases (min_tokens == max_tokens, min_tokens == 0)
+5) Multiple stop conditions
+"""
+
+import pytest
+
+from vllm import LLM, SamplingParams
+from vllm.outputs import RequestOutput
+
+# Test configuration
+TEST_MODEL = "facebook/opt-125m"  # Small model for fast CI execution
+GREEDY = 0.0  # Deterministic generation for consistent testing
+
+
+class MinTokensTestCase:
+    """Data class for min_tokens test scenarios"""
+
+    def __init__(
+        self,
+        name: str,
+        min_tokens: int,
+        max_tokens: int,
+        stop: str | list[str] | None = None,
+        expected_min_len: int | None = None,
+        expected_exact_len: int | None = None,
+    ):
+        self.name = name
+        self.min_tokens = min_tokens
+        self.max_tokens = max_tokens
+        self.stop = stop
+        self.expected_min_len = expected_min_len or min_tokens
+        self.expected_exact_len = expected_exact_len
+
+    def __str__(self):
+        return (
+            f"{self.name}: min={self.min_tokens}, "
+            f"max={self.max_tokens}, stop={self.stop}"
+        )
+
+
+# Test scenarios covering all critical cases
+MIN_TOKENS_TEST_CASES = [
+    # === BASIC FUNCTIONALITY (should work) ===
+    MinTokensTestCase(
+        name="basic_min_tokens_no_stop",
+        min_tokens=8,
+        max_tokens=20,
+        stop=None,
+        expected_min_len=8,
+    ),
+    MinTokensTestCase(
+        name="min_tokens_zero",
+        min_tokens=0,
+        max_tokens=10,
+        stop=None,
+        expected_min_len=0,
+    ),
+    MinTokensTestCase(
+        name="min_equals_max_no_stop",
+        min_tokens=15,
+        max_tokens=15,
+        stop=None,
+        expected_exact_len=15,
+    ),
+    # === STOP STRINGS WITH MIN_TOKENS ===
+    # These tests expose the detokenizer bug where stop strings
+    # bypass min_tokens
+    # Using mathematically guaranteed approach with wide stop nets
+    pytest.param(
+        MinTokensTestCase(
+            name="min_tokens_with_comprehensive_stops",
+            min_tokens=5,
+            max_tokens=20,
+            stop=[
+                "a",
+                "e",
+                "i",
+                "o",
+                "u",
+                "t",
+                "n",
+                "s",
+                "r",
+                "l",
+                " ",
+            ],
+            expected_min_len=5,
+        ),
+        marks=pytest.mark.xfail(
+            reason=(
+                "Known bug #21987: stop strings bypass min_tokens (fixed by PR #22014)"
+            ),
+            strict=False,
+        ),
+        id="min_tokens_with_comprehensive_stops",
+    ),
+    pytest.param(
+        MinTokensTestCase(
+            name="min_tokens_with_simple_char_stop",
+            min_tokens=3,
+            max_tokens=15,
+            stop=["e", "a", " "],
+            expected_min_len=3,
+        ),
+        marks=pytest.mark.xfail(
+            reason=(
+                "Known bug #21987: stop strings bypass min_tokens (fixed by PR #22014)"
+            ),
+            strict=False,
+        ),
+        id="min_tokens_with_simple_char_stop",
+    ),
+    # === EOS TOKEN WITH MIN_TOKENS (potential LogitsProcessor bug) ===
+    # These test the MinTokensLogitsProcessor handling of EOS tokens
+    pytest.param(
+        MinTokensTestCase(
+            name="min_equals_max_eos_only",
+            min_tokens=20,
+            max_tokens=20,
+            stop=None,  # Relies on default EOS token behavior
+            expected_exact_len=20,
+        ),
+        marks=pytest.mark.xfail(
+            reason=("Potential logits-processor bug: EOS tokens may bypass min_tokens"),
+            strict=False,
+        ),
+        id="min_equals_max_eos_only",
+    ),
+    # === EDGE CASES ===
+    MinTokensTestCase(
+        name="large_min_tokens",
+        min_tokens=50,
+        max_tokens=60,
+        stop=None,
+        expected_min_len=50,
+    ),
+    MinTokensTestCase(
+        name="min_tokens_with_empty_stop_list",
+        min_tokens=5,
+        max_tokens=15,
+        stop=[],  # Empty stop list
+        expected_min_len=5,
+    ),
+]
+
+
+@pytest.fixture(scope="module")
+def llm_v1():
+    """Create V1 LLM instance for testing"""
+    llm = LLM(
+        model=TEST_MODEL,
+        tensor_parallel_size=1,
+        max_model_len=1024,  # Small context for fast testing
+        enforce_eager=True,  # Avoid graph compilation overhead
+    )
+    return llm
+
+
+def get_token_count(output: RequestOutput) -> int:
+    """Extract token count from LLM output"""
+    if not output.outputs:
+        return 0
+    return len(output.outputs[0].token_ids)
+
+
+def assert_min_tokens_satisfied(
+    output: RequestOutput, test_case: MinTokensTestCase
+) -> None:
+    """Assert that min_tokens requirement is satisfied"""
+    token_count = get_token_count(output)
+    stop_reason = output.outputs[0].stop_reason if output.outputs else "no output"
+
+    if test_case.expected_exact_len is not None:
+        # Exact length requirement
+        assert token_count == test_case.expected_exact_len, (
+            f"Expected exactly {test_case.expected_exact_len} tokens, "
+            f"got {token_count} tokens. "
+            f"Stop reason: {stop_reason}"
+        )
+    else:
+        # Minimum length requirement
+        assert token_count >= (test_case.expected_min_len or 0), (
+            f"Expected at least {test_case.expected_min_len} tokens, "
+            f"got {token_count} tokens. "
+            f"Stop reason: {stop_reason}"
+        )
+
+
+@pytest.mark.parametrize(
+    "test_case",
+    MIN_TOKENS_TEST_CASES,
+    ids=lambda tc: tc.name,
+)
+def test_min_tokens_comprehensive(llm_v1: LLM, test_case: MinTokensTestCase):
+    """
+    Comprehensive test for min_tokens functionality in V1 engine.
+
+    This test covers all critical scenarios for min_tokens:
+    - Basic functionality (should work)
+    - Stop strings with min_tokens (known bug)
+    - EOS tokens with min_tokens (potential bug)
+    - Edge cases
+
+    Args:
+        llm_v1: V1 LLM instance
+        test_case: Test scenario parameters
+    """
+    # Known failing cases are handled via param-level xfail marks above.
+
+    # Create sampling parameters
+    sampling_params = SamplingParams(
+        min_tokens=test_case.min_tokens,
+        max_tokens=test_case.max_tokens,
+        stop=test_case.stop,
+        temperature=GREEDY,
+        include_stop_str_in_output=True,  # Include stop strings for debugging
+    )
+
+    # Use simple prompt. Comprehensive stop lists should catch any generation
+    prompt = "Hello"
+
+    # Generate output
+    outputs = llm_v1.generate([prompt], sampling_params)
+
+    assert len(outputs) == 1, "Expected exactly one output"
+    output = outputs[0]
+
+    # Debug information
+    token_count = get_token_count(output)
+    generated_text = output.outputs[0].text if output.outputs else ""
+    stop_reason = output.outputs[0].stop_reason if output.outputs else "unknown"
+
+    print(f"\nTest: {test_case.name}")
+    print(f"Generated {token_count} tokens")
+    print(f"Stop reason: {stop_reason}")
+    print(f"Generated text: {repr(generated_text)}")
+    print(f"Expected min: {test_case.expected_min_len}")
+    if test_case.expected_exact_len:
+        print(f"Expected exact: {test_case.expected_exact_len}")
+
+    # Validate min_tokens requirement
+    assert_min_tokens_satisfied(output, test_case)
+
+
+def test_min_tokens_basic_functionality(llm_v1: LLM):
+    """
+    Test basic min_tokens functionality without stop conditions.
+
+    This is a baseline test that should always pass and validates
+    that min_tokens works correctly in the simple case.
+    """
+    sampling_params = SamplingParams(min_tokens=10, max_tokens=20, temperature=GREEDY)
+
+    prompt = "Once upon a time"
+    outputs = llm_v1.generate([prompt], sampling_params)
+
+    assert len(outputs) == 1
+    token_count = get_token_count(outputs[0])
+
+    assert token_count >= 10, f"Expected at least 10 tokens, got {token_count}"
+    assert token_count <= 20, f"Expected at most 20 tokens, got {token_count}"
+
+
+@pytest.mark.xfail(
+    reason=("Known bug #21987: stop strings bypass min_tokens (fixed by PR #22014)"),
+    strict=False,
+)
+def test_min_tokens_stop_strings_bug(llm_v1: LLM):
+    """
+    Test the specific bug where stop strings bypass min_tokens.
+
+    This test specifically reproduces the bug Calvin is fixing in PR #22014.
+    It should fail until that fix is merged.
+
+    Strategy: Use guaranteed stop characters that will appear
+    in any generated text.
+    """
+    # If the bug is fixed upstream, this test will XPASS
+
+    sampling_params = SamplingParams(
+        min_tokens=15,
+        max_tokens=50,
+        # Common letter; likely appears early
+        stop=["e"],
+        temperature=GREEDY,
+        include_stop_str_in_output=True,
+    )
+
+    # Simple prompt that will generate text containing "e"
+    prompt = "The quick brown fox"
+    outputs = llm_v1.generate([prompt], sampling_params)
+
+    assert len(outputs) == 1
+    token_count = get_token_count(outputs[0])
+    generated_text = outputs[0].outputs[0].text if outputs[0].outputs else ""
+
+    # Debug info to understand what happened
+    print(f"Generated text: {repr(generated_text)}")
+    print(f"Token count: {token_count}")
+    print(f"Contains 'e': {'e' in generated_text}")
+
+    # This assertion should fail due to the bug - if stop string is found early,
+    # the model should still continue generating until min_tokens is reached
+    stop_reason = (
+        outputs[0].outputs[0].stop_reason if outputs[0].outputs else "no output"
+    )
+    assert token_count >= 15, (
+        "Bug confirmed: "
+        f"{token_count} tokens < min_tokens=15. "
+        f"Reason: {stop_reason}. "
+        f"Text: {repr(generated_text)}"
+    )
+
+
+@pytest.mark.xfail(
+    reason=("Known bug #21987: stop strings bypass min_tokens (fixed by PR #22014)"),
+    strict=False,
+)
+def test_min_tokens_stop_strings_guaranteed_early_trigger(llm_v1: LLM):
+    """
+    Guaranteed test for stop strings bypassing min_tokens bug.
+
+    Strategy: Use very low temperature and multiple common stop strings
+    to virtually guarantee early detection, combined with long min_tokens
+    to ensure the bug is exposed regardless of model behavior.
+    """
+    # If the bug is fixed upstream, this test will XPASS
+
+    sampling_params = SamplingParams(
+        min_tokens=50,  # Set high min_tokens to ensure bug detection
+        max_tokens=200,
+        # Use multiple very common patterns - at least one will appear
+        stop=["e", "a", "i", "o", "u", " ", "t", "n", "s", "r"],
+        temperature=GREEDY,
+        include_stop_str_in_output=True,
+    )
+
+    # Simple prompt that will generate some text
+    prompt = "The cat"
+    outputs = llm_v1.generate([prompt], sampling_params)
+
+    assert len(outputs) == 1
+    token_count = get_token_count(outputs[0])
+    generated_text = outputs[0].outputs[0].text if outputs[0].outputs else ""
+    stop_reason = outputs[0].outputs[0].stop_reason if outputs[0].outputs else "unknown"
+
+    print(f"Generated text: {repr(generated_text)}")
+    print(f"Token count: {token_count}")
+    print(f"Stop reason: {stop_reason}")
+
+    # With the bug, this will fail because ANY of the common characters
+    # will trigger early termination before min_tokens=50 is reached
+    # It's virtually impossible to generate 50 tokens without hitting
+    # at least one of: e, a, i, o, u, space, t, n, s, r
+    finish_reason = (
+        outputs[0].outputs[0].finish_reason if outputs[0].outputs else "unknown"
+    )
+
+    print(f"Finish reason: {finish_reason}")
+
+    if finish_reason == "stop":
+        assert token_count >= 50, (
+            "Bug confirmed: "
+            f"{token_count} tokens < min_tokens=50. "
+            f"Reason: {finish_reason}. "
+            f"Text: {repr(generated_text)}"
+        )
+
+
+@pytest.mark.xfail(
+    reason=("Potential logits-processor bug: EOS tokens may bypass min_tokens"),
+    strict=False,
+)
+def test_min_tokens_eos_behavior(llm_v1: LLM):
+    """
+    Verify EOS handling with and without min_tokens.
+
+    - Without min_tokens: expect early EOS -> finish_reason == "stop",
+      stop_reason is None, and generated tokens < max_tokens (25).
+    - With min_tokens: EOS should be blocked until min_tokens is reached
+      (finish_reason == "length"); verify that eos_token_id does not appear
+      in generated token_ids.
+    """
+    # tokenizer + eos id
+    tokenizer = llm_v1.get_tokenizer()
+    eos_token_id = tokenizer.eos_token_id
+
+    prompt = "Give a file extension."
+    max_toks = 32
+
+    # Case 1: WITHOUT min_tokens
+    sp_no_min = SamplingParams(
+        max_tokens=max_toks,
+        temperature=GREEDY,
+    )
+    out_no_min = llm_v1.generate([prompt], sp_no_min)
+    assert len(out_no_min) == 1
+    choice_no_min = out_no_min[0].outputs[0]
+
+    ids_no_min = choice_no_min.token_ids or []
+    finish_no_min = choice_no_min.finish_reason
+    stop_no_min = choice_no_min.stop_reason
+
+    print(
+        "[no-min] tokens=",
+        len(ids_no_min),
+        " finish=",
+        finish_no_min,
+        " stop_reason=",
+        stop_no_min,
+    )
+
+    assert finish_no_min == "stop", (
+        f"Expected finish_reason 'stop' without min_tokens, got {finish_no_min}"
+    )
+    assert stop_no_min is None, (
+        "For EOS-based stop (no user stop strings), stop_reason should be None."
+    )
+    assert len(ids_no_min) < max_toks, (
+        f"Expected early EOS with < {max_toks} tokens, got {len(ids_no_min)}"
+    )
+
+    # Case 2: WITH min_tokens
+    sp_with_min = SamplingParams(
+        min_tokens=max_toks,
+        max_tokens=max_toks,
+        temperature=GREEDY,
+    )
+    out_with_min = llm_v1.generate([prompt], sp_with_min)
+    assert len(out_with_min) == 1
+    choice_with_min = out_with_min[0].outputs[0]
+
+    ids_with_min = choice_with_min.token_ids or []
+    finish_with_min = choice_with_min.finish_reason
+    stop_with_min = choice_with_min.stop_reason
+
+    print(
+        "[with-min] tokens=",
+        len(ids_with_min),
+        " finish=",
+        finish_with_min,
+        " stop_reason=",
+        stop_with_min,
+    )
+
+    # Exact length reached; EOS should have been blocked
+    assert len(ids_with_min) == max_toks, (
+        f"Expected exactly {max_toks} tokens with min_tokens; got {len(ids_with_min)}"
+    )
+    assert finish_with_min == "length", (
+        f"Expected finish_reason 'length'; got {finish_with_min}"
+    )
+    assert eos_token_id not in ids_with_min, (
+        "EOS token id should not appear when min_tokens prevents early EOS."
+    )
+
+
+def test_min_tokens_validation():
+    """
+    Test that SamplingParams correctly validates min_tokens parameters.
+
+    This tests the parameter validation logic in SamplingParams.
+    """
+    # Valid cases
+    SamplingParams(min_tokens=0, max_tokens=10)
+    SamplingParams(min_tokens=5, max_tokens=10)
+    SamplingParams(min_tokens=10, max_tokens=10)
+
+    # Invalid cases
+    with pytest.raises(
+        ValueError,
+        match="min_tokens must be greater than or equal to 0",
+    ):
+        SamplingParams(min_tokens=-1, max_tokens=10)
+
+    with pytest.raises(
+        ValueError,
+        match="min_tokens must be less than or equal to max_tokens",
+    ):
+        SamplingParams(min_tokens=15, max_tokens=10)
+
+
+if __name__ == "__main__":
+    """
+    Run tests locally for development.
+    
+    Usage:
+        cd vllm/
+        python -m pytest tests/v1/e2e/general/test_min_tokens.py -v
+    """
+    pytest.main([__file__, "-v"])

third_party/vllm/tests/v1/e2e/general/test_pooling_chunked_prefill.py

@@ -0,0 +1,168 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+import pytest
+import torch.nn as nn
+
+from vllm.platforms import current_platform
+
+prompt = """
+Generals gathered in their masses
+Just like witches at black masses
+Evil minds that plot destruction
+Sorcerer of death's construction
+In the fields, the bodies burning
+As the war machine keeps turning
+Death and hatred to mankind
+Poisoning their brainwashed minds
+Oh, Lord, yeah
+
+Politicians hide themselves away
+They only started the war
+Why should they go out to fight?
+They leave that all to the poor, yeah
+Time will tell on their power minds
+Making war just for fun
+Treating people just like pawns in chess
+Wait till their judgment day comes, yeah
+
+Now, in darkness, world stops turning
+Ashes where their bodies burning
+No more war pigs have the power
+Hand of God has struck the hour
+Day of Judgment, God is calling
+On their knees, the war pigs crawling
+Begging mercies for their sins
+Satan, laughing, spreads his wings
+Oh, Lord, yeah
+"""
+
+
+class WrapperPooler(nn.Module):
+    def __init__(self, pooler):
+        super().__init__()
+        self.pooler = pooler
+        self.chunks = []
+
+    def get_pooling_updates(self, task):
+        return self.pooler.get_pooling_updates(task)
+
+    def forward(
+        self,
+        hidden_states,
+        pooling_metadata,
+    ):
+        self.chunks.append(hidden_states.shape[0])
+        return self.pooler(hidden_states, pooling_metadata)
+
+
+def inject_pooler(self):
+    model = self.get_model()
+    wrapper = WrapperPooler(model.pooler)
+    model.pooler = wrapper
+
+
+def retrieve_chunks(self):
+    model = self.get_model()
+    chunks = model.pooler.chunks
+    model.pooler.chunks = []
+    return chunks
+
+
+@pytest.mark.skipif(not current_platform.is_cuda(), reason="CUDA not available")
+def test_pooling_chunked_prefill(vllm_runner, monkeypatch):
+    """Test chunked prefill for pooling models with LastPool."""
+
+    with monkeypatch.context() as m:
+        m.setenv("VLLM_ALLOW_INSECURE_SERIALIZATION", "1")
+        model_id = "Qwen/Qwen3-Embedding-0.6B"
+
+        chunk_size = 10
+
+        # Set chunking parameters to force chunked prefill
+        # Note: Chunked prefill is automatically handled by vLLM
+        # internally based on the model size and prompt
+        with vllm_runner(
+            model_id,
+            runner="pooling",
+            long_prefill_token_threshold=chunk_size,
+            tensor_parallel_size=1,
+            enforce_eager=True,
+            enable_chunked_prefill=True,
+        ) as llm:
+            llm.get_llm().llm_engine.collective_rpc(inject_pooler)
+
+            tokenizer = llm.get_llm().get_tokenizer()
+            tokens = tokenizer(prompt)["input_ids"]
+            prompt_len = len(tokens)
+            full_chunks, last_chunk = divmod(prompt_len, chunk_size)
+            expected_chunks = [chunk_size] * full_chunks
+            if last_chunk:
+                expected_chunks.append(last_chunk)
+            llm.embed([prompt])
+            chunks = llm.get_llm().llm_engine.collective_rpc(retrieve_chunks)[0]
+
+        # Check that PoolerWrapper was called and chunks were received
+        assert len(chunks) > 1
+        assert chunks == expected_chunks
+
+        # Disable chunked prefill
+        with vllm_runner(
+            model_id,
+            runner="pooling",
+            tensor_parallel_size=1,
+            enforce_eager=True,
+        ) as llm:
+            llm.get_llm().llm_engine.collective_rpc(inject_pooler)
+            llm.embed([prompt])
+            chunks = llm.get_llm().llm_engine.collective_rpc(retrieve_chunks)[0]
+
+        # Check that PoolerWrapper was called and no chunks were received
+        assert len(chunks) == 1
+        assert chunks[0] == prompt_len
+
+
+@pytest.mark.skipif(not current_platform.is_cuda(), reason="CUDA not available")
+def test_pooling_prefix_cache(vllm_runner, monkeypatch):
+    """Test chunked prefill for pooling models with LastPool."""
+
+    verses = prompt.split("\n\n")
+
+    with monkeypatch.context() as m:
+        m.setenv("VLLM_ALLOW_INSECURE_SERIALIZATION", "1")
+        model_id = "Qwen/Qwen3-Embedding-0.6B"
+
+        with vllm_runner(
+            model_id,
+            runner="pooling",
+            enable_prefix_caching=True,
+            tensor_parallel_size=1,
+            enforce_eager=True,
+        ) as llm:
+            llm.get_llm().llm_engine.collective_rpc(inject_pooler)
+            tokenizer = llm.get_llm().get_tokenizer()
+
+            prompt1 = "\n\n".join([verses[0], verses[1]])
+            prompt2 = "\n\n".join([verses[0], verses[2]])
+            tokens1 = tokenizer(prompt1)["input_ids"]
+            tokens2 = tokenizer(prompt2)["input_ids"]
+            prompt1_len = len(tokens1)
+            prompt2_len = len(tokens2)
+
+            llm.embed([prompt1])
+            chunks = llm.get_llm().llm_engine.collective_rpc(retrieve_chunks)[0]
+
+            assert len(chunks) == 1
+            assert chunks[0] == prompt1_len
+
+            llm.embed([prompt2])
+            chunks = llm.get_llm().llm_engine.collective_rpc(retrieve_chunks)[0]
+
+            assert len(chunks) == 1
+            assert chunks[0] <= prompt1_len
+            assert chunks[0] < prompt2_len
+
+            vllm_config = llm.get_llm().llm_engine.vllm_config
+            cache_config = vllm_config.cache_config
+            print(f"{cache_config=}")
+            # Prefixes are cached in blocks
+            assert (prompt2_len - chunks[0]) % cache_config.block_size == 0

third_party/vllm/tests/v1/e2e/general/test_streaming_input.py

@@ -0,0 +1,657 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""
+End-to-end tests for the streaming input feature in AsyncLLM.
+
+These tests verify that:
+1. Streaming inputs work correctly with bunched inputs (queued)
+2. Streaming inputs work correctly with spaced out inputs
+3. Outputs are equivalent whether inputs are bunched or spaced
+4. Cancelling the output stream correctly aborts the session
+5. Closing the input stream correctly signals completion
+6. Queued inputs are cancelled when the session is aborted
+"""
+
+import asyncio
+from collections.abc import AsyncGenerator
+
+import pytest
+import pytest_asyncio
+
+from vllm import SamplingParams
+from vllm.engine.protocol import StreamingInput
+from vllm.outputs import RequestOutput
+from vllm.platforms import current_platform
+from vllm.sampling_params import RequestOutputKind
+from vllm.utils.torch_utils import set_default_torch_num_threads
+from vllm.v1.engine.async_llm import AsyncLLM
+
+if not current_platform.is_cuda():
+    pytest.skip(reason="V1 currently only supported on CUDA.", allow_module_level=True)
+
+# Use a small model that doesn't require authentication for fast tests
+MODEL = "facebook/opt-125m"
+
+
+@pytest_asyncio.fixture(scope="module", loop_scope="module")
+async def engine():
+    """Create an AsyncLLM engine for the test.
+
+    Note: Using function scope because pytest_asyncio creates a new event loop
+    for each test, and the output_handler task gets cancelled between tests
+    with module scope.
+    """
+    from vllm.engine.arg_utils import AsyncEngineArgs
+
+    engine_args = AsyncEngineArgs(
+        model=MODEL, enforce_eager=True, gpu_memory_utilization=0.7
+    )
+    with set_default_torch_num_threads(1):
+        engine = AsyncLLM.from_engine_args(engine_args)
+    try:
+        yield engine
+    finally:
+        engine.shutdown()
+        await asyncio.sleep(0.1)
+
+
+def get_sampling_params(max_tokens: int = 20) -> SamplingParams:
+    """Create sampling params for streaming input tests."""
+    return SamplingParams(
+        max_tokens=max_tokens,
+        ignore_eos=True,
+        output_kind=RequestOutputKind.DELTA,
+        temperature=0.0,  # Deterministic for reproducibility
+    )
+
+
+async def collect_outputs(
+    output_gen: AsyncGenerator[RequestOutput, None],
+) -> tuple[list[RequestOutput], str]:
+    """Collect all outputs from a generate call, return outputs and full text."""
+    outputs: list[RequestOutput] = []
+    full_text = ""
+    async for output in output_gen:
+        outputs.append(output)
+        if output.outputs and output.outputs[0].text:
+            full_text += output.outputs[0].text
+    return outputs, full_text
+
+
+@pytest.mark.asyncio(loop_scope="module")
+async def test_streaming_input_bunched(engine: AsyncLLM):
+    """Test streaming input where all inputs are sent at once (bunched/queued).
+
+    This tests the case where multiple inputs arrive before any completes.
+    The inputs should be queued and processed in sequence.
+    """
+    request_id = "test_bunched"
+    sampling_params = get_sampling_params(max_tokens=10)
+
+    # Create an input generator that yields all inputs quickly
+    async def bunched_input_generator() -> AsyncGenerator[StreamingInput, None]:
+        # Send multiple inputs rapidly - they should be queued
+        yield StreamingInput(prompt="Hello, my name is")
+        yield StreamingInput(prompt=" Alice and I like")
+        yield StreamingInput(prompt=" to code in Python")
+
+    outputs, full_text = await collect_outputs(
+        engine.generate(
+            bunched_input_generator(),
+            sampling_params,
+            request_id,
+        )
+    )
+
+    # Verify we got outputs
+    assert len(outputs) > 0, "Should have received outputs"
+
+    # Verify the final output is marked as finished
+    assert outputs[-1].finished, "Last output should be marked as finished"
+
+    # Verify intermediate outputs are not marked as finished
+    for output in outputs[:-1]:
+        assert not output.finished, "Intermediate outputs should not be finished"
+
+    # Verify we generated some text
+    assert len(full_text) > 0, "Should have generated text"
+    print(f"Bunched test generated: {full_text}")
+
+
+@pytest.mark.asyncio(loop_scope="module")
+async def test_streaming_input_spaced(engine: AsyncLLM):
+    """Test streaming input where inputs are spaced out.
+
+    This tests the case where each input completes processing before the
+    next one is sent. Each chunk should be prefilled, generate tokens,
+    then the next chunk should be processed.
+    """
+    request_id = "test_spaced"
+    sampling_params = get_sampling_params(max_tokens=10)
+
+    # Track when each input is sent
+    input_times: list[float] = []
+    outputs_per_chunk: list[int] = [0, 0, 0]
+    current_chunk = 0
+
+    async def spaced_input_generator() -> AsyncGenerator[StreamingInput, None]:
+        nonlocal current_chunk
+        import time
+
+        # First input
+        input_times.append(time.time())
+        yield StreamingInput(prompt="Hello, my name is")
+        current_chunk = 0
+
+        # Wait for some outputs to be generated
+        await asyncio.sleep(0.5)
+
+        # Second input
+        input_times.append(time.time())
+        current_chunk = 1
+        yield StreamingInput(prompt=" Alice and I like")
+
+        # Wait for some outputs
+        await asyncio.sleep(0.5)
+
+        # Third input
+        input_times.append(time.time())
+        current_chunk = 2
+        yield StreamingInput(prompt=" to code in Python")
+
+    outputs: list[RequestOutput] = []
+    full_text = ""
+
+    async for output in engine.generate(
+        spaced_input_generator(),
+        sampling_params,
+        request_id,
+    ):
+        outputs.append(output)
+        if output.outputs and output.outputs[0].text:
+            full_text += output.outputs[0].text
+            outputs_per_chunk[current_chunk] += 1
+
+    # Verify we got outputs
+    assert len(outputs) > 0, "Should have received outputs"
+
+    # Verify the final output is marked as finished
+    assert outputs[-1].finished, "Last output should be marked as finished"
+
+    # Verify we received outputs from multiple chunks
+    # (with spaced inputs, we should see outputs distributed across chunks)
+    chunks_with_outputs = sum(1 for c in outputs_per_chunk if c > 0)
+    assert chunks_with_outputs >= 1, "Should have outputs from at least one chunk"
+
+    print(f"Spaced test generated: {full_text}")
+    print(f"Outputs per chunk: {outputs_per_chunk}")
+
+
+@pytest.mark.asyncio(loop_scope="module")
+async def test_streaming_input_output_equivalence(engine: AsyncLLM):
+    """Test that bunched and spaced inputs produce equivalent outputs.
+
+    When the same prompts are provided either bunched or spaced,
+    the final concatenated output should be the same (with deterministic
+    sampling).
+    """
+    prompts = ["Hello, my name is", " Bob and I work", " at Anthropic"]
+    sampling_params = get_sampling_params(max_tokens=15)
+
+    # Test bunched inputs
+    async def bunched_gen() -> AsyncGenerator[StreamingInput, None]:
+        for prompt in prompts:
+            yield StreamingInput(prompt=prompt)
+
+    _, bunched_text = await collect_outputs(
+        engine.generate(bunched_gen(), sampling_params, "equiv_bunched")
+    )
+
+    # Test spaced inputs (same prompts, but with delays)
+    async def spaced_gen() -> AsyncGenerator[StreamingInput, None]:
+        for prompt in prompts:
+            yield StreamingInput(prompt=prompt)
+            await asyncio.sleep(0.3)
+
+    _, spaced_text = await collect_outputs(
+        engine.generate(spaced_gen(), sampling_params, "equiv_spaced")
+    )
+
+    # Both should produce the same output since we use temperature=0
+    assert bunched_text == spaced_text, (
+        f"Bunched and spaced should produce same output.\n"
+        f"Bunched: {bunched_text!r}\n"
+        f"Spaced: {spaced_text!r}"
+    )
+
+    print(f"Equivalence test passed. Generated: {bunched_text}")
+
+
+@pytest.mark.asyncio(loop_scope="module")
+async def test_streaming_input_cancel_output_stream(engine: AsyncLLM):
+    """Test that cancelling the output stream aborts the entire session.
+
+    When the consumer cancels iteration over the output generator,
+    the session should be aborted including any queued inputs.
+    """
+    request_id = "test_cancel_output"
+    sampling_params = get_sampling_params(max_tokens=1000)
+
+    input_completed = asyncio.Event()
+    input_task_cancelled = False
+
+    async def slow_input_generator() -> AsyncGenerator[StreamingInput, None]:
+        nonlocal input_task_cancelled
+        try:
+            yield StreamingInput(prompt="Tell me a very long story about")
+            yield StreamingInput(prompt=" a dragon and a knight")
+
+            # This should be cancelled before we get here
+            await asyncio.sleep(10)
+            yield StreamingInput(prompt=" who become friends")
+            input_completed.set()
+        except asyncio.CancelledError:
+            input_task_cancelled = True
+            raise
+
+    outputs_received = 0
+    output_gen = engine.generate(slow_input_generator(), sampling_params, request_id)
+
+    # Collect a few outputs then cancel
+    try:
+        async for output in output_gen:
+            outputs_received += 1
+            if outputs_received >= 5:
+                # Cancel by breaking out of the loop (generator will be GC'd)
+                break
+    finally:
+        # Explicitly close the generator to ensure cleanup
+        await output_gen.aclose()
+
+    # Give time for cleanup
+    await asyncio.sleep(0.5)
+
+    # Verify we got some outputs before cancelling
+    assert outputs_received >= 5, "Should have received outputs before cancel"
+
+    # Verify the input task was cancelled
+    assert input_task_cancelled, "Input task should have been cancelled"
+
+    # Verify the session is properly cleaned up
+    assert not engine.output_processor.has_unfinished_requests(), (
+        "Should have no unfinished requests after cancel"
+    )
+
+    print(f"Cancel test passed. Received {outputs_received} outputs before cancel")
+
+
+@pytest.mark.asyncio(loop_scope="module")
+async def test_streaming_input_close_signals_completion(engine: AsyncLLM):
+    """Test that closing the input stream signals completion.
+
+    When the input generator finishes (naturally or via return),
+    the session should complete with finished=True on the last output.
+    """
+    request_id = "test_close_completion"
+    sampling_params = get_sampling_params(max_tokens=15)
+
+    input_generator_finished = False
+
+    async def limited_input_generator() -> AsyncGenerator[StreamingInput, None]:
+        nonlocal input_generator_finished
+        yield StreamingInput(prompt="What is 2 + 2? The answer is")
+        # Generator finishes naturally here
+        input_generator_finished = True
+
+    outputs, _ = await collect_outputs(
+        engine.generate(limited_input_generator(), sampling_params, request_id)
+    )
+
+    # Verify the input generator completed
+    assert input_generator_finished, "Input generator should have finished"
+
+    # Verify we got a finished output
+    assert len(outputs) > 0, "Should have received outputs"
+    assert outputs[-1].finished, "Last output should be marked as finished"
+
+    # Verify the session is cleaned up
+    assert not engine.output_processor.has_unfinished_requests(), (
+        "Should have no unfinished requests"
+    )
+
+    print("Close completion test passed")
+
+
+@pytest.mark.asyncio(loop_scope="module")
+async def test_streaming_input_abort_queued_inputs(engine: AsyncLLM):
+    """Test that aborting the session cancels queued inputs.
+
+    When multiple inputs are queued and the session is aborted,
+    all pending inputs should be cancelled.
+    """
+    request_id = "test_abort_queued"
+    # Use large max_tokens to ensure we have time to queue inputs
+    sampling_params = get_sampling_params(max_tokens=2000)
+
+    inputs_sent = 0
+    input_cancelled = False
+
+    async def many_inputs_generator() -> AsyncGenerator[StreamingInput, None]:
+        nonlocal inputs_sent, input_cancelled
+        try:
+            # Send several inputs to fill the queue
+            for i in range(10):
+                yield StreamingInput(prompt=f" Part {i}: Tell me about the number {i}.")
+                inputs_sent += 1
+                # Small delay to interleave with output processing
+                await asyncio.sleep(0.05)
+        except asyncio.CancelledError:
+            input_cancelled = True
+            raise
+
+    outputs_received = 0
+    output_gen = engine.generate(many_inputs_generator(), sampling_params, request_id)
+
+    try:
+        async for output in output_gen:
+            outputs_received += 1
+            # Cancel after receiving some outputs
+            if outputs_received >= 10:
+                break
+    finally:
+        await output_gen.aclose()
+
+    # Give time for cleanup
+    await asyncio.sleep(0.5)
+
+    # Verify we received some outputs
+    assert outputs_received >= 10, "Should have received outputs before abort"
+
+    # Verify the input generator was cancelled OR finished naturally
+    # (it might finish naturally if all inputs were sent before cancel)
+    assert input_cancelled or inputs_sent == 10, (
+        f"Input generator should have been cancelled or completed. "
+        f"cancelled={input_cancelled}, inputs_sent={inputs_sent}"
+    )
+
+    # Verify the session is cleaned up
+    assert not engine.output_processor.has_unfinished_requests(), (
+        "Should have no unfinished requests after abort"
+    )
+
+    print(
+        f"Abort queued test passed. Sent {inputs_sent} inputs, "
+        f"received {outputs_received} outputs"
+    )
+
+
+@pytest.mark.asyncio(loop_scope="module")
+async def test_streaming_input_error_propagation(engine: AsyncLLM):
+    """Test that errors in the input generator are propagated to the caller."""
+    request_id = "test_error_propagation"
+    sampling_params = get_sampling_params(max_tokens=20)
+
+    class InputError(Exception):
+        pass
+
+    async def error_input_generator() -> AsyncGenerator[StreamingInput, None]:
+        yield StreamingInput(prompt="Start with this")
+        await asyncio.sleep(0.1)
+        raise InputError("Simulated input error")
+
+    # Note: The current implementation catches exceptions and puts them
+    # in the queue, so we should get the error when iterating outputs
+    with pytest.raises(InputError, match="Simulated input error"):
+        async for _ in engine.generate(
+            error_input_generator(), sampling_params, request_id
+        ):
+            pass
+
+    # Give time for cleanup
+    await asyncio.sleep(0.3)
+
+    # Verify the session is cleaned up
+    assert not engine.output_processor.has_unfinished_requests(), (
+        "Should have no unfinished requests after error"
+    )
+
+
+@pytest.mark.asyncio(loop_scope="module")
+async def test_streaming_input_multiple_concurrent_sessions(engine: AsyncLLM):
+    """Test multiple concurrent streaming input sessions.
+
+    Multiple streaming sessions should be able to run concurrently
+    without interfering with each other.
+    """
+    num_sessions = 3
+    results: list[tuple[str, str]] = []
+
+    async def run_session(session_id: int) -> tuple[str, str]:
+        request_id = f"test_concurrent_{session_id}"
+        sampling_params = get_sampling_params(max_tokens=10)
+
+        prompts = [f"Session {session_id}: Hello", f" world from session {session_id}"]
+
+        async def input_gen() -> AsyncGenerator[StreamingInput, None]:
+            for prompt in prompts:
+                yield StreamingInput(prompt=prompt)
+                await asyncio.sleep(0.1)
+
+        _, text = await collect_outputs(
+            engine.generate(input_gen(), sampling_params, request_id)
+        )
+        return request_id, text
+
+    # Run sessions concurrently
+    tasks = [asyncio.create_task(run_session(i)) for i in range(num_sessions)]
+    results = await asyncio.gather(*tasks)
+
+    # Verify all sessions completed
+    assert len(results) == num_sessions
+
+    for request_id, text in results:
+        assert len(text) > 0, f"Session {request_id} should have generated text"
+        print(f"{request_id}: {text}")
+
+    # Verify cleanup
+    assert not engine.output_processor.has_unfinished_requests()
+
+
+@pytest.mark.asyncio(loop_scope="module")
+async def test_streaming_input_per_chunk_sampling_params(engine: AsyncLLM):
+    """Test that per-chunk sampling params are respected.
+
+    Each StreamingInput can have its own sampling_params.
+    """
+    request_id = "test_per_chunk_params"
+    base_params = get_sampling_params(max_tokens=10)
+
+    async def variable_params_generator() -> AsyncGenerator[StreamingInput, None]:
+        # First chunk with base params
+        yield StreamingInput(prompt="Count to five:", sampling_params=base_params)
+
+        # Second chunk with different max_tokens
+        chunk_params = get_sampling_params(max_tokens=5)
+        yield StreamingInput(
+            prompt=" Now count backwards:", sampling_params=chunk_params
+        )
+
+    outputs, full_text = await collect_outputs(
+        engine.generate(variable_params_generator(), base_params, request_id)
+    )
+
+    assert len(outputs) > 0, "Should have received outputs"
+    assert outputs[-1].finished, "Last output should be finished"
+    assert len(full_text) > 0, "Should have generated text"
+
+    print(f"Per-chunk params test generated: {full_text}")
+
+
+@pytest.mark.asyncio(loop_scope="module")
+async def test_streaming_input_empty_generator(engine: AsyncLLM):
+    """Test behavior when the input generator yields nothing.
+
+    An empty generator should still produce a finished output.
+    """
+    request_id = "test_empty_generator"
+    sampling_params = get_sampling_params(max_tokens=10)
+
+    async def empty_generator() -> AsyncGenerator[StreamingInput, None]:
+        # Don't yield anything
+        return
+        yield  # Make it a generator
+
+    outputs: list[RequestOutput] = []
+    async for output in engine.generate(empty_generator(), sampling_params, request_id):
+        outputs.append(output)
+
+    # Should still get a finished marker
+    assert len(outputs) >= 1, "Should receive at least one output"
+    assert outputs[-1].finished, "Should have a finished output"
+
+
+@pytest.mark.asyncio(loop_scope="module")
+async def test_streaming_input_single_chunk(engine: AsyncLLM):
+    """Test streaming input with a single chunk.
+
+    This is effectively the same as a regular non-streaming request,
+    but using the streaming input API.
+    """
+    request_id = "test_single_chunk"
+    sampling_params = get_sampling_params(max_tokens=15)
+
+    async def single_chunk_generator() -> AsyncGenerator[StreamingInput, None]:
+        yield StreamingInput(prompt="What color is the sky? The sky is")
+
+    outputs, full_text = await collect_outputs(
+        engine.generate(single_chunk_generator(), sampling_params, request_id)
+    )
+
+    assert len(outputs) > 0
+    assert outputs[-1].finished
+    assert "blue" in full_text.lower() or len(full_text) > 0
+
+    print(f"Single chunk test generated: {full_text}")
+
+
+@pytest.mark.asyncio(loop_scope="module")
+async def test_streaming_input_reuse_request_id(engine: AsyncLLM):
+    """Test that request IDs can be reused after a session completes."""
+    request_id = "test_reuse_id"
+    sampling_params = get_sampling_params(max_tokens=5)
+
+    # First session
+    async def gen1() -> AsyncGenerator[StreamingInput, None]:
+        yield StreamingInput(prompt="First session")
+
+    _, text1 = await collect_outputs(
+        engine.generate(gen1(), sampling_params, request_id)
+    )
+
+    # Second session with same ID
+    async def gen2() -> AsyncGenerator[StreamingInput, None]:
+        yield StreamingInput(prompt="Second session")
+
+    _, text2 = await collect_outputs(
+        engine.generate(gen2(), sampling_params, request_id)
+    )
+
+    assert len(text1) > 0
+    assert len(text2) > 0
+    assert not engine.output_processor.has_unfinished_requests()
+
+    print(f"Reuse ID test: session 1: {text1}, session 2: {text2}")
+
+
+@pytest.mark.asyncio(loop_scope="module")
+async def test_streaming_input_validation_errors(engine: AsyncLLM):
+    """Test that invalid configurations raise appropriate errors."""
+
+    async def dummy_generator() -> AsyncGenerator[StreamingInput, None]:
+        yield StreamingInput(prompt="test")
+
+    # Test n > 1 is rejected
+    with pytest.raises(ValueError, match="Input streaming not currently supported"):
+        params_n2 = SamplingParams(max_tokens=10, n=2)
+        async for _ in engine.generate(dummy_generator(), params_n2, "test_n2"):
+            pass
+
+    # Test FINAL_ONLY is rejected
+    with pytest.raises(ValueError, match="Input streaming not currently supported"):
+        params_final = SamplingParams(
+            max_tokens=10, output_kind=RequestOutputKind.FINAL_ONLY
+        )
+        async for _ in engine.generate(dummy_generator(), params_final, "test_final"):
+            pass
+
+    # Test stop strings are rejected
+    with pytest.raises(ValueError, match="Input streaming not currently supported"):
+        params_stop = SamplingParams(max_tokens=10, stop=["stop"])
+        async for _ in engine.generate(dummy_generator(), params_stop, "test_stop"):
+            pass
+
+
+@pytest.mark.asyncio(loop_scope="module")
+async def test_streaming_input_delayed_generator_exit(engine: AsyncLLM):
+    """Test that output generator exits when input generator closes after outputs.
+
+    This tests the case where:
+    1. Multiple inputs are sent and fully processed
+    2. The engine has finished
+    3. The input generator doesn't exit until after the engine finishes
+    4. The output generator should exit properly once the input generator exits
+    """
+    request_id = "test_delayed_exit"
+    sampling_params = get_sampling_params(max_tokens=10)
+
+    engine_finished_event = asyncio.Event()
+    input_generator_exited = False
+    finish_count = 0
+
+    async def delayed_exit_input_generator() -> AsyncGenerator[StreamingInput, None]:
+        nonlocal input_generator_exited
+        # Send all inputs immediately
+        yield StreamingInput(prompt="Hello, my name is")
+        yield StreamingInput(prompt=" Alice")
+
+        # Wait until the engine has finished generating before exiting
+        await engine_finished_event.wait()
+
+        # Add a small delay to ensure we're testing the "delayed exit" case
+        await asyncio.sleep(0.1)
+        input_generator_exited = True
+
+    outputs: list[RequestOutput] = []
+    full_text = ""
+
+    async for output in engine.generate(
+        delayed_exit_input_generator(), sampling_params, request_id
+    ):
+        outputs.append(output)
+        if output.outputs and output.outputs[0].text:
+            full_text += output.outputs[0].text
+
+        # Signal when the engine finishes both input chunks (each gets a finish_reason)
+        # Note: output.finished will be False while input stream is open
+        if output.outputs and output.outputs[0].finish_reason is not None:
+            finish_count += 1
+            if finish_count == 2:
+                engine_finished_event.set()
+
+    # Verify the input generator exited properly
+    assert input_generator_exited, (
+        "Input generator should have exited after engine finished"
+    )
+
+    # Verify we got outputs
+    assert len(outputs) > 0, "Should have received outputs"
+
+    # Verify we generated some text
+    assert len(full_text) > 0, "Should have generated text"
+
+    # Verify the session is cleaned up
+    assert not engine.output_processor.has_unfinished_requests(), (
+        "Should have no unfinished requests"
+    )
+
+    print(f"Delayed exit test passed. Generated: {full_text}")

third_party/vllm/tests/v1/e2e/spec_decode/test_async_spec_decode.py

@@ -0,0 +1,131 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""
+Test that verifies no implicit GPU-CPU synchronization occurs during
+speculative decoding generation under expected conditions.
+"""
+
+import multiprocessing
+import sys
+import traceback
+
+import pytest
+import torch
+
+
+@pytest.fixture
+def sync_tracker():
+    """
+    Fixture that patches CommonAttentionMetadata.seq_lens_cpu to detect
+    lazy init syncs. Prints stack traces immediately when syncs occur.
+    """
+    from vllm.v1.attention.backend import CommonAttentionMetadata
+
+    # Shared counter for cross-process communication (inherited by fork)
+    sync_count = multiprocessing.Value("i", 0)
+
+    # Save original property
+    original_prop = CommonAttentionMetadata.seq_lens_cpu
+    original_fget = original_prop.fget
+
+    # Create tracking wrapper
+    def tracking_seq_lens_cpu(self):
+        if self._seq_lens_cpu is None:
+            # Increment counter
+            with sync_count.get_lock():
+                sync_count.value += 1
+                count = sync_count.value
+            # Print stack trace immediately (shows in subprocess output)
+            print(f"\n{'=' * 60}", file=sys.stderr)
+            print(f"SYNC #{count}: seq_lens_cpu lazy init triggered!", file=sys.stderr)
+            print(f"{'=' * 60}", file=sys.stderr)
+            traceback.print_stack(file=sys.stderr)
+            print(f"{'=' * 60}\n", file=sys.stderr)
+            sys.stderr.flush()
+        return original_fget(self)
+
+    # Apply patch
+    CommonAttentionMetadata.seq_lens_cpu = property(tracking_seq_lens_cpu)
+
+    class SyncTracker:
+        @property
+        def count(self) -> int:
+            return sync_count.value
+
+        def assert_no_sync(self, msg: str = ""):
+            count = sync_count.value
+            assert count == 0, (
+                f"Unexpected GPU-CPU sync: seq_lens_cpu lazy init triggered "
+                f"{count} times. See stack traces above. {msg}"
+            )
+
+    yield SyncTracker()
+
+    # Restore original property
+    CommonAttentionMetadata.seq_lens_cpu = original_prop
+    torch._dynamo.reset()
+
+
+# Test configurations: (model, spec_model, method, num_spec_tokens, backend_env)
+SPEC_DECODE_CONFIGS = [
+    pytest.param(
+        "meta-llama/Llama-3.2-1B-Instruct",
+        "nm-testing/Llama3_2_1B_speculator.eagle3",
+        "eagle3",
+        2,
+        id="eagle3-llama",
+    ),
+    pytest.param(
+        "eagle618/deepseek-v3-random",
+        "eagle618/eagle-deepseek-v3-random",
+        "eagle",
+        2,
+        id="eagle-mla-deepseek",
+    ),
+]
+
+
+@pytest.mark.parametrize(
+    "model,spec_model,method,num_spec_tokens",
+    SPEC_DECODE_CONFIGS,
+)
+def test_no_sync_with_spec_decode(
+    sync_tracker,
+    model: str,
+    spec_model: str,
+    method: str,
+    num_spec_tokens: int,
+):
+    """
+    Test that no implicit GPU-CPU sync occurs during speculative decoding
+    generation.
+    """
+    # Import vLLM AFTER sync_tracker fixture has applied the patch
+    from vllm import LLM, SamplingParams
+    from vllm.distributed import cleanup_dist_env_and_memory
+
+    llm = LLM(
+        model=model,
+        max_model_len=256,
+        speculative_config={
+            "method": method,
+            "num_speculative_tokens": num_spec_tokens,
+            "model": spec_model,
+        },
+        enforce_eager=True,
+        async_scheduling=True,
+    )
+
+    outputs = llm.generate(
+        ["Hello, my name is"],
+        SamplingParams(temperature=0, max_tokens=10),
+    )
+
+    assert len(outputs) == 1
+    assert len(outputs[0].outputs[0].text) > 0
+
+    del llm
+    torch.accelerator.empty_cache()
+    cleanup_dist_env_and_memory()
+
+    sync_tracker.assert_no_sync()

third_party/vllm/tests/v1/e2e/spec_decode/test_lora_with_spec_decode.py

@@ -0,0 +1,139 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""
+This script contains:
+1. test lora with speculative decoding for batch inference
+"""
+
+import random
+
+import numpy as np
+import pytest
+import torch
+
+from vllm import LLM, SamplingParams
+from vllm.distributed import cleanup_dist_env_and_memory
+from vllm.lora.request import LoRARequest
+from vllm.platforms import current_platform
+
+LORA_TEST_PROMPT_MAP: dict[str, str] = {}
+
+LORA_TEST_PROMPT_MAP["premjatin/qwen-linear-algebra-coder"] = """
+### INSTRUCTION:
+You are an AI assistant that generates Python code to solve linear
+algebra problems.
+
+### PROBLEM:
+Find the eigenvalues and eigenvectors of the following 3x3 matrix:
+[[3, 2, 0],
+ [2, 3, 0],
+ [0, 0, 2]]
+
+### OUTPUT FORMAT (STRICT):
+Numbers should be represented as integers only.
+
+### PYTHON SOLUTION:
+"""
+
+SEED = 42
+
+
+@pytest.mark.skipif(not current_platform.is_cuda(), reason="CUDA not available")
+@pytest.mark.parametrize(
+    "model_setup",
+    [
+        (
+            "eagle3",
+            "Qwen/Qwen3-1.7B",
+            "AngelSlim/Qwen3-1.7B_eagle3",
+            "premjatin/qwen-linear-algebra-coder",
+            1,
+        )
+    ],
+)
+def test_batch_inference_correctness(
+    monkeypatch: pytest.MonkeyPatch,
+    model_setup: tuple[str, str, str, str, int],
+):
+    """
+    Compare the outputs of a LLM with only Lora and a LLM with both SD and Lora.
+    Should be the same and no failure when doing batch inference.
+    model_setup: (method, model_name, spec_model_name, lora_path, tp_size)
+    """
+    with monkeypatch.context() as m:
+        # Disable randomness
+        m.setenv("CUBLAS_WORKSPACE_CONFIG", ":4096:8")
+        torch.manual_seed(SEED)
+        np.random.seed(SEED)
+        random.seed(SEED)
+        torch.cuda.manual_seed_all(SEED)
+        torch.backends.cudnn.benchmark = False
+        torch.backends.cudnn.deterministic = True
+
+        method, model_name, spec_model_name, lora_path, tp_size = model_setup
+
+        # without speculative decoding
+        ref_llm = LLM(
+            model=model_name,
+            trust_remote_code=True,
+            tensor_parallel_size=tp_size,
+            max_model_len=2048,
+            max_num_seqs=4,
+            enable_lora=True,
+            max_loras=1,
+            max_cpu_loras=1,
+            max_lora_rank=16,
+        )
+
+        prompts = [LORA_TEST_PROMPT_MAP[lora_path]] * 100
+        lora_request = LoRARequest("adapter", 1, lora_path)
+        sampling_params = SamplingParams(
+            temperature=0.0, top_p=1.0, top_k=-1, seed=SEED, max_tokens=128
+        )
+
+        ref_outputs = ref_llm.generate(
+            prompts, sampling_params, lora_request=lora_request
+        )
+        del ref_llm
+        torch.accelerator.empty_cache()
+        cleanup_dist_env_and_memory()
+
+        lora_spec_llm = LLM(
+            model=model_name,
+            trust_remote_code=True,
+            tensor_parallel_size=tp_size,
+            speculative_config={
+                "method": method,
+                "model": spec_model_name,
+                "num_speculative_tokens": 3,
+                "max_model_len": 2048,
+            },
+            max_model_len=2048,
+            max_num_seqs=4,
+            enable_lora=True,
+            max_loras=1,
+            max_cpu_loras=1,
+            max_lora_rank=16,
+        )
+
+        lora_spec_outputs = lora_spec_llm.generate(
+            prompts, sampling_params, lora_request=lora_request
+        )
+
+        matches = 0
+        misses = 0
+        for ref_output, spec_output in zip(ref_outputs, lora_spec_outputs):
+            if ref_output.outputs[0].text == spec_output.outputs[0].text:
+                matches += 1
+            else:
+                misses += 1
+                print(f"ref_output: {ref_output.outputs[0].text}")
+                print(f"spec_output: {spec_output.outputs[0].text}")
+
+        # Heuristic: expect at least 90% of the prompts to match exactly
+        # Upon failure, inspect the outputs to check for inaccuracy.
+        print(f"match ratio: {matches}/{len(ref_outputs)}")
+        assert matches > int(0.90 * len(ref_outputs))
+        del lora_spec_llm
+        torch.accelerator.empty_cache()
+        cleanup_dist_env_and_memory()