wehub-resource-sync
351 lines
12 KiB
Python
351 lines
12 KiB
Python
# SPDX-License-Identifier: Apache-2.0
|
|
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
|
|
|
"""E2E tests for ``thinking_token_budget`` with reasoning models.
|
|
|
|
Covers Qwen3-0.6B and Qwen3.5 FP8 + MTP.
|
|
"""
|
|
|
|
import asyncio
|
|
import json
|
|
from typing import Literal
|
|
|
|
import openai
|
|
import pytest
|
|
import pytest_asyncio
|
|
|
|
from tests.utils import RemoteOpenAIServer, multi_gpu_only, requires_fp8
|
|
from vllm.platforms import current_platform
|
|
from vllm.tokenizers import get_tokenizer
|
|
|
|
MODEL_NAME = "Qwen/Qwen3-0.6B"
|
|
QWEN35_FP8_MTP_MODEL = "Qwen/Qwen3.5-35B-A3B-FP8"
|
|
MESSAGES = [{"role": "user", "content": "What is 1+1? Be concise."}]
|
|
THINK_BUDGET = 5
|
|
|
|
REASONING_START_STR = "<think>"
|
|
REASONING_END_STR = "</think>"
|
|
|
|
|
|
def _count_reasoning_decode_token_ids_between_markers(
|
|
full_token_ids: list[int],
|
|
reasoning_start_ids: list[int],
|
|
reasoning_end_ids: list[int],
|
|
) -> int | None:
|
|
"""Count decode tokens in the thinking span (after last start, before first end)."""
|
|
|
|
if not reasoning_start_ids or not reasoning_end_ids:
|
|
raise ValueError("reasoning marker token id lists must be non-empty")
|
|
|
|
def _last_subseq_index(haystack: list[int], needle: list[int]) -> int:
|
|
n = len(needle)
|
|
if n > len(haystack):
|
|
return -1
|
|
for i in range(len(haystack) - n, -1, -1):
|
|
if haystack[i : i + n] == needle:
|
|
return i
|
|
return -1
|
|
|
|
last_start = _last_subseq_index(full_token_ids, reasoning_start_ids)
|
|
if last_start < 0:
|
|
return None
|
|
|
|
pos_after_start = last_start + len(reasoning_start_ids)
|
|
end_n = len(reasoning_end_ids)
|
|
for j in range(pos_after_start, len(full_token_ids) - end_n + 1):
|
|
if full_token_ids[j : j + end_n] == reasoning_end_ids:
|
|
return j - pos_after_start
|
|
return len(full_token_ids) - pos_after_start
|
|
|
|
|
|
@pytest.fixture(scope="module")
|
|
def server():
|
|
args = [
|
|
"--reasoning-parser",
|
|
"qwen3",
|
|
"--reasoning-config",
|
|
'{"reasoning_start_str": "<think>", "reasoning_end_str": "</think>"}',
|
|
"--max-model-len",
|
|
"2048",
|
|
"--enforce-eager",
|
|
"--gpu-memory-utilization",
|
|
"0.4",
|
|
"--no-async-scheduling",
|
|
]
|
|
# thinking_token_budget is not yet supported by the V2 model runner.
|
|
env_dict = {"VLLM_USE_V2_MODEL_RUNNER": "0"}
|
|
with RemoteOpenAIServer(MODEL_NAME, args, env_dict=env_dict) as remote_server:
|
|
yield remote_server
|
|
|
|
|
|
@pytest.fixture(scope="module")
|
|
def server_with_auto_reasoning_config():
|
|
args = [
|
|
"--reasoning-parser",
|
|
"qwen3",
|
|
"--max-model-len",
|
|
"2048",
|
|
"--enforce-eager",
|
|
"--gpu-memory-utilization",
|
|
"0.4",
|
|
"--no-async-scheduling",
|
|
]
|
|
# thinking_token_budget is not yet supported by the V2 model runner.
|
|
env_dict = {"VLLM_USE_V2_MODEL_RUNNER": "0"}
|
|
with RemoteOpenAIServer(MODEL_NAME, args, env_dict=env_dict) as remote_server:
|
|
yield remote_server
|
|
|
|
|
|
@pytest.fixture(scope="module")
|
|
def server_qwen35_fp8_mtp_tp2():
|
|
"""Qwen3.5-35B FP8 with MTP speculative decoding and tensor parallel size 2."""
|
|
if current_platform.device_count() < 2:
|
|
pytest.skip("Need at least 2 GPUs for --tensor-parallel-size 2")
|
|
if not current_platform.supports_fp8():
|
|
pytest.skip("FP8 is not supported on this platform")
|
|
|
|
spec_cfg = {
|
|
"method": "mtp",
|
|
"num_speculative_tokens": 2,
|
|
"max_model_len": 32768,
|
|
}
|
|
args = [
|
|
"--tensor-parallel-size",
|
|
"2",
|
|
"--max-model-len",
|
|
"32768",
|
|
"--speculative-config",
|
|
json.dumps(spec_cfg),
|
|
"--reasoning-parser",
|
|
"qwen3",
|
|
"--reasoning-config",
|
|
json.dumps(
|
|
{
|
|
"reasoning_start_str": REASONING_START_STR,
|
|
"reasoning_end_str": REASONING_END_STR,
|
|
}
|
|
),
|
|
]
|
|
# thinking_token_budget is not yet supported by the V2 model runner.
|
|
env_dict: dict[str, str] = {"VLLM_USE_V2_MODEL_RUNNER": "0"}
|
|
# With 4+ GPUs, run TP=2 on physical devices 2,3 so module-scoped 0.6B servers
|
|
# on 0,1 do not exhaust memory on the same devices as this worker.
|
|
if current_platform.device_count() >= 4:
|
|
env_dict["CUDA_VISIBLE_DEVICES"] = "2,3"
|
|
|
|
with RemoteOpenAIServer(
|
|
QWEN35_FP8_MTP_MODEL,
|
|
args,
|
|
max_wait_seconds=3000,
|
|
env_dict=env_dict,
|
|
) as remote_server:
|
|
yield remote_server
|
|
|
|
|
|
@pytest_asyncio.fixture
|
|
async def client(request, server, server_with_auto_reasoning_config):
|
|
server_map = {
|
|
"default": server,
|
|
"auto_config": server_with_auto_reasoning_config,
|
|
}
|
|
target_server = server_map[request.param]
|
|
async with target_server.get_async_client() as async_client:
|
|
yield async_client
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
@pytest.mark.parametrize("client", ["default", "auto_config"], indirect=True)
|
|
async def test_thinking_token_budget_mixed_requests(client: openai.AsyncOpenAI):
|
|
"""Test that mixed requests (some with thinking_token_budget, some without)
|
|
complete successfully without errors."""
|
|
|
|
response_with_budget = await client.chat.completions.create(
|
|
model=MODEL_NAME,
|
|
messages=MESSAGES,
|
|
max_tokens=100,
|
|
extra_body={"thinking_token_budget": THINK_BUDGET},
|
|
)
|
|
response_without_budget = await client.chat.completions.create(
|
|
model=MODEL_NAME,
|
|
messages=MESSAGES,
|
|
max_tokens=100,
|
|
)
|
|
|
|
msg_with = response_with_budget.choices[0].message
|
|
msg_without = response_without_budget.choices[0].message
|
|
|
|
assert msg_with.content or getattr(msg_with, "reasoning", None)
|
|
assert msg_without.content or getattr(msg_without, "reasoning", None)
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
@pytest.mark.parametrize("client", ["default", "auto_config"], indirect=True)
|
|
async def test_thinking_token_budget_limits_reasoning(client: openai.AsyncOpenAI):
|
|
"""Test that thinking_token_budget limits the number of reasoning tokens.
|
|
|
|
Counts reasoning decode tokens by id, which is robust to how tokens are
|
|
grouped into streamed chunks (a single chunk can carry several tokens under
|
|
async scheduling / stream_interval > 1). Counting chunks under-counts.
|
|
"""
|
|
|
|
tokenizer = get_tokenizer(tokenizer_name=MODEL_NAME)
|
|
start_ids = list(tokenizer.encode(REASONING_START_STR, add_special_tokens=False))
|
|
end_ids = list(tokenizer.encode(REASONING_END_STR, add_special_tokens=False))
|
|
|
|
prompt_token_ids: list[int] = []
|
|
decode_token_ids: list[int] = []
|
|
stream = await client.chat.completions.create(
|
|
model=MODEL_NAME,
|
|
messages=MESSAGES,
|
|
max_tokens=100,
|
|
stream=True,
|
|
extra_body={"thinking_token_budget": THINK_BUDGET, "return_token_ids": True},
|
|
)
|
|
async for chunk in stream:
|
|
if not chunk.choices:
|
|
continue
|
|
if getattr(chunk, "prompt_token_ids", None):
|
|
prompt_token_ids = list(chunk.prompt_token_ids)
|
|
delta_ids = getattr(chunk.choices[0], "token_ids", None)
|
|
if delta_ids:
|
|
decode_token_ids.extend(delta_ids)
|
|
|
|
reasoning_token_count = _count_reasoning_decode_token_ids_between_markers(
|
|
prompt_token_ids + decode_token_ids, start_ids, end_ids
|
|
)
|
|
assert reasoning_token_count is not None, "missing reasoning start marker in ids"
|
|
assert reasoning_token_count == THINK_BUDGET, (
|
|
f"reasoning tokens ({reasoning_token_count}) != "
|
|
f"thinking_token_budget ({THINK_BUDGET})"
|
|
)
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
@multi_gpu_only(num_gpus=2)
|
|
@requires_fp8
|
|
async def test_thinking_token_budget_qwen35_fp8_mtp_concurrent_mixed_budget_and_plain(
|
|
server_qwen35_fp8_mtp_tp2,
|
|
):
|
|
"""Concurrent chat requests: some with ``thinking_token_budget``, some without.
|
|
|
|
Exercises the scheduler / input processor under a mixed batch on the same
|
|
Qwen3.5 FP8 + MTP (TP=2) server. Budgeted calls are checked with
|
|
``_count_reasoning_decode_token_ids_between_markers`` on full token ids.
|
|
"""
|
|
|
|
_batch_spec: list[tuple[Literal["budget"], int] | tuple[Literal["plain"], None]] = [
|
|
("budget", 1),
|
|
("budget", 12),
|
|
("plain", None),
|
|
("budget", 20),
|
|
("budget", 14),
|
|
("plain", None),
|
|
("plain", None),
|
|
("budget", 12),
|
|
("plain", None),
|
|
]
|
|
|
|
tokenizer = get_tokenizer(tokenizer_name=QWEN35_FP8_MTP_MODEL)
|
|
start_ids = list(tokenizer.encode(REASONING_START_STR, add_special_tokens=False))
|
|
end_ids = list(tokenizer.encode(REASONING_END_STR, add_special_tokens=False))
|
|
|
|
async with server_qwen35_fp8_mtp_tp2.get_async_client() as client:
|
|
|
|
async def budgeted_call(expected_budget: int):
|
|
return await client.chat.completions.create(
|
|
model=QWEN35_FP8_MTP_MODEL,
|
|
messages=MESSAGES,
|
|
max_tokens=256,
|
|
stream=False,
|
|
extra_body={
|
|
"thinking_token_budget": expected_budget,
|
|
"return_token_ids": True,
|
|
},
|
|
)
|
|
|
|
async def plain_call():
|
|
return await client.chat.completions.create(
|
|
model=QWEN35_FP8_MTP_MODEL,
|
|
messages=MESSAGES,
|
|
max_tokens=256,
|
|
stream=False,
|
|
)
|
|
|
|
coros = []
|
|
for row in _batch_spec:
|
|
if row[0] == "budget":
|
|
b = row[1]
|
|
assert isinstance(b, int)
|
|
coros.append(budgeted_call(b))
|
|
else:
|
|
coros.append(plain_call())
|
|
results = await asyncio.gather(*coros)
|
|
|
|
for i, (response, (kind, expected_budget)) in enumerate(
|
|
zip(results, _batch_spec, strict=True)
|
|
):
|
|
msg = response.choices[0].message
|
|
assert msg.content or getattr(msg, "reasoning", None), (
|
|
f"index {i} ({kind}): empty message"
|
|
)
|
|
|
|
if kind == "budget":
|
|
assert expected_budget is not None
|
|
assert response.prompt_token_ids is not None
|
|
assert response.choices[0].token_ids is not None
|
|
full_ids = list(response.prompt_token_ids) + list(
|
|
response.choices[0].token_ids
|
|
)
|
|
n_reason = _count_reasoning_decode_token_ids_between_markers(
|
|
full_ids, start_ids, end_ids
|
|
)
|
|
assert n_reason is not None, f"index {i}: missing reasoning start in ids"
|
|
assert n_reason == expected_budget, (
|
|
f"index {i}: reasoning decode token ids ({n_reason}) != "
|
|
f"thinking_token_budget ({expected_budget})"
|
|
)
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
@pytest.mark.parametrize("client", ["default", "auto_config"], indirect=True)
|
|
async def test_streaming_with_thinking_disabled_stays_in_content(
|
|
client: openai.AsyncOpenAI,
|
|
):
|
|
request_kwargs = {
|
|
"model": MODEL_NAME,
|
|
"messages": [
|
|
{
|
|
"role": "user",
|
|
"content": "Which is larger, 4 or 12?"
|
|
" Output exactly one token: 4 or 12.",
|
|
}
|
|
],
|
|
"max_tokens": 16,
|
|
"temperature": 0.0,
|
|
"extra_body": {"chat_template_kwargs": {"enable_thinking": False}},
|
|
}
|
|
|
|
response = await client.chat.completions.create(**request_kwargs)
|
|
message = response.choices[0].message
|
|
assert message.content is not None and message.content.strip() != ""
|
|
assert getattr(message, "reasoning", None) in (None, "")
|
|
|
|
stream = await client.chat.completions.create(
|
|
**request_kwargs,
|
|
stream=True,
|
|
)
|
|
|
|
content_chunks = []
|
|
reasoning_chunks = []
|
|
async for chunk in stream:
|
|
if not chunk.choices:
|
|
continue
|
|
delta = chunk.choices[0].delta
|
|
if getattr(delta, "content", None):
|
|
content_chunks.append(delta.content)
|
|
if getattr(delta, "reasoning", None):
|
|
reasoning_chunks.append(delta.reasoning)
|
|
|
|
assert "".join(content_chunks).strip() != ""
|
|
assert reasoning_chunks == []
|