[Bugfix] Fix diffusion response fanout and reusable request parameters

tests/entrypoints/openai/chat_completion/test_serving_chat.py

@@ -20,6 +20,7 @@
 from tests.utils import RemoteOpenAIServer
 from vllm._aiter_ops import is_aiter_found_and_supported
 from vllm.config import MultiModalConfig
+from vllm.config.diffusion import DiffusionConfig
 from vllm.entrypoints.generate.base.protocol import (
     DeltaMessage,
     RequestResponseMetadata,
@@ -54,6 +55,7 @@
 from vllm.tokenizers import get_tokenizer
 from vllm.tokenizers.mistral import MistralTokenizer
 from vllm.tokenizers.registry import cached_tokenizer_from_config
+from vllm.utils.diffusion import validate_diffusion_sampling_params
 from vllm.v1.engine.async_llm import AsyncLLM
 from vllm.v1.metrics.stats import RequestStateStats
 
@@ -689,6 +691,80 @@ async def _stream_request_outputs(
         yield request_output
 
 
+@pytest.mark.asyncio
+@pytest.mark.parametrize("diffusion_samples", [False, True])
+async def test_streaming_diffusion_samples_preserves_all_choices(diffusion_samples):
+    """All fanout children must reach the client, just as ordinary n does."""
+    request = ChatCompletionRequest(
+        model="test-model",
+        messages=[{"role": "user", "content": "Test"}],
+        stream=True,
+        stream_options={"include_usage": True, "continuous_usage_stats": True},
+        n=1 if diffusion_samples else 4,
+        vllm_xargs={
+            "diffusion_seed_canvas": [0, 1],
+            "diffusion_pinned": [0],
+            "diffusion_samples": 4,
+        }
+        if diffusion_samples
+        else None,
+    )
+    params = request.to_sampling_params(2, {})
+    validate_diffusion_sampling_params(
+        params, canvas_length=2, vocab_size=128, async_scheduling=True
+    )
+    output = _make_metrics_request_output()
+    output.outputs = [
+        CompletionOutput(
+            index=i,
+            text=f"sample {i}",
+            token_ids=[100, 101],
+            cumulative_logprob=None,
+            logprobs=None,
+            finish_reason="length",
+        )
+        for i in range(params.n)
+    ]
+    serving = _build_minimal_metrics_serving_chat(True)
+    serving.model_config = MockModelConfig()
+    serving.model_config.is_diffusion = True
+    serving.diffusion_config = DiffusionConfig()
+    chunks = [
+        json.loads(line.removeprefix("data: "))
+        async for line in serving.chat_completion_stream_generator(
+            request,
+            _single_request_output(output),
+            "chatcmpl-fanout",
+            "test-model",
+            conversation=[{"role": "user", "content": "Test"}],
+            tokenizer=None,
+            request_metadata=RequestResponseMetadata(request_id="chatcmpl-fanout"),
+        )
+        if line.strip() != "data: [DONE]"
+    ]
+    assert not [chunk for chunk in chunks if "error" in chunk], chunks
+    finished = [
+        choice["index"]
+        for chunk in chunks
+        for choice in chunk["choices"]
+        if choice.get("finish_reason") is not None
+    ]
+    assert finished == list(range(4))
+    assert chunks[-1]["usage"]["completion_tokens"] == 8
+    assert not chunks[-1].get("metrics")
+    full = await serving.chat_completion_full_generator(
+        request,
+        _single_request_output(output),
+        "chatcmpl-fanout",
+        "test-model",
+        [{"role": "user", "content": "Test"}],
+        None,
+        RequestResponseMetadata(request_id="chatcmpl-fanout"),
+    )
+    assert len(full.choices) == 4
+    assert full.metrics is None
+
+
 async def _collect_metrics_stream_chunks(
     serving: OpenAIServingChat,
     request: ChatCompletionRequest,
@@ -1330,7 +1406,12 @@ async def return_model_name(*args, **kwargs):
 
 @pytest.mark.asyncio
 @pytest.mark.parametrize("stream", [True, False])
-async def test_admission_rejection_escapes_before_response_starts(stream):
+@pytest.mark.parametrize(
+    "samples,is_diffusion,expected_n", [(None, False, 1), (4, False, 1), (4, True, 4)]
+)
+async def test_admission_rejection_escapes_before_response_starts(
+    stream, samples, is_diffusion, expected_n
+):
     """Overload rejections must propagate out of create_chat_completion.
 
     For streaming this is the only chance to return a real HTTP status: once
@@ -1340,21 +1421,31 @@ async def test_admission_rejection_escapes_before_response_starts(stream):
     mock_engine = MagicMock(spec=AsyncLLM)
     mock_engine.errored = False
     mock_engine.model_config = MockModelConfig()
+    mock_engine.model_config.is_diffusion = is_diffusion
     mock_engine.input_processor = MagicMock()
     mock_engine.renderer = _build_renderer(mock_engine.model_config)
     mock_engine.check_admission.side_effect = QueueOverflowError()
 
     serving_chat = _build_serving_chat(mock_engine)
+    serving_chat.diffusion_config = DiffusionConfig(max_samples=4)
     req = ChatCompletionRequest(
         model=MODEL_NAME,
         messages=[{"role": "user", "content": "what is 1+1?"}],
         stream=stream,
+        vllm_xargs={
+            "diffusion_samples": samples,
+            "diffusion_seed_canvas": [0, 1],
+            "diffusion_pinned": [0],
+        }
+        if samples is not None
+        else None,
     )
 
     with pytest.raises(QueueOverflowError):
         await serving_chat.create_chat_completion(req)
 
     mock_engine.generate.assert_not_called()
+    mock_engine.check_admission.assert_called_once_with(expected_n)
 
 
 @pytest.mark.asyncio

tests/entrypoints/openai/completion/test_completion_error.py

@@ -1,12 +1,14 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 
+import json
 from dataclasses import dataclass, field
 from typing import Any
 from unittest.mock import AsyncMock, MagicMock
 
 import pytest
 
+from vllm.config.diffusion import DiffusionConfig
 from vllm.config.multimodal import MultiModalConfig
 from vllm.entrypoints.generate.base.protocol import RequestResponseMetadata
 from vllm.entrypoints.openai.completion.protocol import CompletionRequest
@@ -21,6 +23,7 @@
 from vllm.renderers.hf import HfRenderer
 from vllm.renderers.online_renderer import OnlineRenderer
 from vllm.tokenizers.registry import cached_tokenizer_from_config
+from vllm.utils.diffusion import validate_diffusion_sampling_params
 from vllm.v1.engine.async_llm import AsyncLLM
 from vllm.v1.metrics.stats import RequestSpecDecodeMetrics, RequestStateStats
 
@@ -149,6 +152,83 @@ def _build_renderer(model_config: MockModelConfig):
     )
 
 
+@pytest.mark.asyncio
+@pytest.mark.parametrize("diffusion_samples", [False, True])
+@pytest.mark.parametrize("num_prompts", [1, 2])
+async def test_streaming_diffusion_samples_preserves_all_choices(
+    diffusion_samples, num_prompts
+):
+    """Fanout choices need distinct streaming slots across batched prompts."""
+    request = CompletionRequest(
+        model=MODEL_NAME,
+        prompt=["Test"] * num_prompts,
+        max_tokens=2,
+        stream=True,
+        stream_options={"include_usage": True},
+        n=1 if diffusion_samples else 4,
+        vllm_xargs={
+            "diffusion_seed_canvas": [0, 1],
+            "diffusion_pinned": [0],
+            "diffusion_samples": 4,
+        }
+        if diffusion_samples
+        else None,
+    )
+    params = request.to_sampling_params(2)
+    validate_diffusion_sampling_params(
+        params, canvas_length=2, vocab_size=128, async_scheduling=True
+    )
+    output = _make_metrics_request_output()
+    output.outputs = [
+        CompletionOutput(
+            index=i,
+            text=f"sample {i}",
+            token_ids=[100, 101],
+            cumulative_logprob=None,
+            logprobs=None,
+            finish_reason="length",
+        )
+        for i in range(params.n)
+    ]
+
+    async def results():
+        for prompt_idx in range(num_prompts):
+            yield prompt_idx, output
+
+    serving = _build_minimal_metrics_serving_completion(True)
+    serving.enable_force_include_usage = False
+    serving.model_config = MockModelConfig()
+    serving.model_config.is_diffusion = True
+    serving.diffusion_config = DiffusionConfig()
+    metadata = RequestResponseMetadata(request_id="cmpl-fanout")
+    chunks = [
+        json.loads(line.removeprefix("data: "))
+        async for line in serving.completion_stream_generator(
+            request,
+            [],
+            results(),
+            "cmpl-fanout",
+            0,
+            MODEL_NAME,
+            num_prompts,
+            None,
+            metadata,
+        )
+        if line.strip() != "data: [DONE]"
+    ]
+    assert not [chunk for chunk in chunks if "error" in chunk], chunks
+    assert [choice["index"] for chunk in chunks for choice in chunk["choices"]] == list(
+        range(4 * num_prompts)
+    )
+    assert chunks[-1]["usage"]["completion_tokens"] == 8 * num_prompts
+    assert not chunks[-1].get("metrics")
+    full = serving.request_output_to_completion_response(
+        [output] * num_prompts, request, "cmpl-fanout", 0, MODEL_NAME, None, metadata
+    )
+    assert len(full.choices) == 4 * num_prompts
+    assert full.metrics is None
+
+
 def test_completion_per_request_metrics_follow_server_flag():
     request = CompletionRequest(model=MODEL_NAME, prompt="Test prompt", max_tokens=10)
     request_output = _make_metrics_request_output()

tests/test_request_input_bounds.py

@@ -20,6 +20,7 @@
 
 import vllm.envs as envs
 from vllm import SamplingParams
+from vllm.config.diffusion import DiffusionConfig
 from vllm.entrypoints.openai.chat_completion.protocol import (
     BatchChatCompletionRequest,
     ChatCompletionRequest,
@@ -28,6 +29,7 @@
 from vllm.entrypoints.openai.responses.protocol import ResponsesRequest
 from vllm.exceptions import VLLMValidationError
 from vllm.inputs.engine import embeds_input
+from vllm.pooling_params import PoolingParams
 from vllm.sampling_params import BeamSearchParams
 from vllm.tokenizers import deepseek_v4_encoding, deepseek_v32_encoding
 from vllm.v1.engine.input_processor import InputProcessor
@@ -94,6 +96,68 @@ def test_prompt_embeds_valid_inputs_preserved(prompt_embeds_processor, token_ids
     assert request.prompt_is_token_ids == mask
 
 
+@pytest.mark.parametrize("skip_clone", [False, True])
+def test_diffusion_fanout_params_can_be_reused(prompt_embeds_processor, skip_clone):
+    """Admitting one prompt must not make the next look like conflicting n input."""
+    processor = prompt_embeds_processor
+    processor.model_config.is_diffusion = True
+    processor.diffusion_config = DiffusionConfig(canvas_length=2)
+    processor.vllm_config.scheduler_config.async_scheduling = True
+    params = SamplingParams(
+        seed=7,
+        max_tokens=8,
+        skip_clone=skip_clone,
+        extra_args={
+            "diffusion_seed_canvas": [1, 2],
+            "diffusion_pinned": [0],
+            "diffusion_samples": 4,
+            "diffusion_read_only": True,
+        },
+    )
+    prompt = {"type": "token", "prompt_token_ids": [1, 2, 3]}
+
+    requests = [
+        processor.process_inputs(req_id, prompt, params, ("generate",))
+        for req_id in ("first", "second")
+    ]
+
+    assert params.n == 1
+    assert params.max_tokens == 8
+    assert not params.ignore_eos
+    for request in requests:
+        assert request.sampling_params is not params
+        assert request.sampling_params.n == 4
+        assert request.sampling_params.max_tokens == 2
+        assert request.sampling_params.ignore_eos
+        assert request.sampling_params.seed == 7
+        assert (request.sampling_params.extra_args is params.extra_args) == skip_clone
+
+
+def test_pooling_admission_preserves_deep_copy(prompt_embeds_processor):
+    """Pooler defaults reach the request without sharing its nested parameters."""
+    params = PoolingParams(extra_kwargs={"values": [1]})
+    request = prompt_embeds_processor.process_inputs(
+        "pool", {"type": "token", "prompt_token_ids": [1]}, params, ("plugin",)
+    )
+
+    assert request.sampling_params is None
+    assert request.pooling_params.task == "plugin"
+    assert request.pooling_params.skip_reading_prefix_cache
+    request.pooling_params.extra_kwargs["values"].append(2)
+    assert params.extra_kwargs == {"values": [1]}
+
+
+def test_input_processor_rejects_unknown_params(prompt_embeds_processor):
+    """Unsupported input retains the public type-validation error."""
+    with pytest.raises(TypeError, match="SamplingParams or PoolingParams"):
+        prompt_embeds_processor.process_inputs(
+            "invalid",
+            {"type": "token", "prompt_token_ids": [1]},
+            {},  # type: ignore[arg-type]
+            ("generate",),
+        )
+
+
 @pytest.mark.parametrize("offset", [-1, 0, 2, 3, 4, 2**127])
 def test_routed_experts_offset_validated_before_engine_submission(offset):
     """Reject invalid offsets even when routed-expert output is disabled."""

vllm/entrypoints/generate/base/serving.py

@@ -37,6 +37,7 @@
     extract_trace_headers,
     log_tracing_disabled_warning,
 )
+from vllm.utils.diffusion import resolve_diffusion_num_samples
 from vllm.v1.metrics.stats import RequestStateStats
 
 logger = init_logger(__name__)
@@ -155,6 +156,7 @@ def __init__(
         self.renderer = engine_client.renderer
         self.input_processor = engine_client.input_processor
         vllm_config = getattr(engine_client, "vllm_config", None)
+        self.diffusion_config = getattr(vllm_config, "diffusion_config", None)
         kv_transfer_config = getattr(vllm_config, "kv_transfer_config", None)
         self.has_kv_connector = kv_transfer_config is not None
 
@@ -171,6 +173,26 @@ def __init__(
             # Never fail server startup over the fingerprint.
             self.system_fingerprint = None
 
+    def _get_num_choices(
+        self, request: ChatCompletionRequest | CompletionRequest
+    ) -> int:
+        """Account for diffusion fanout without mutating the client's n."""
+        n = request.n or 1
+        extra = request.vllm_xargs
+        if (
+            not extra
+            or "diffusion_samples" not in extra
+            or not self.model_config.is_diffusion
+        ):
+            return n
+        return resolve_diffusion_num_samples(
+            n,
+            extra,
+            max_samples=self.diffusion_config.max_samples
+            if self.diffusion_config is not None
+            else 32,
+        )
+
     def _preflight(self, n: int = 1) -> None:
         """Engine checks that must run before a response is started.