0
tests/ut/patch/platform/__init__.py
Normal file
0
tests/ut/patch/platform/__init__.py
Normal file
82
tests/ut/patch/platform/test_deepseek_v4_thinking.py
Normal file
82
tests/ut/patch/platform/test_deepseek_v4_thinking.py
Normal file
@@ -0,0 +1,82 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
|
||||
from vllm.entrypoints.openai.chat_completion.protocol import ChatCompletionRequest
|
||||
from vllm.tokenizers import deepseek_v4
|
||||
|
||||
|
||||
class FakeTokenizer:
|
||||
vocab_size = 1
|
||||
|
||||
def get_added_vocab(self):
|
||||
return {}
|
||||
|
||||
def encode(self, text, add_special_tokens=False, **kwargs):
|
||||
return text
|
||||
|
||||
|
||||
def test_deepseek_v4_reasoning_effort_accepts_latest_values():
|
||||
for reasoning_effort in ("none", "minimal", "low", "medium", "high", "xhigh", "max"):
|
||||
request = ChatCompletionRequest(
|
||||
model="deepseek-v4",
|
||||
messages=[{"role": "user", "content": "hi"}],
|
||||
reasoning_effort=reasoning_effort,
|
||||
)
|
||||
assert request.reasoning_effort == reasoning_effort
|
||||
|
||||
|
||||
def test_reasoning_effort_enables_thinking_unless_user_overrides():
|
||||
request = ChatCompletionRequest(
|
||||
model="deepseek-v4",
|
||||
messages=[{"role": "user", "content": "hi"}],
|
||||
reasoning_effort="high",
|
||||
)
|
||||
params = request.build_chat_params(None, "auto")
|
||||
assert params.chat_template_kwargs["enable_thinking"] is True
|
||||
|
||||
request = ChatCompletionRequest(
|
||||
model="deepseek-v4",
|
||||
messages=[{"role": "user", "content": "hi"}],
|
||||
reasoning_effort="none",
|
||||
)
|
||||
params = request.build_chat_params(None, "auto")
|
||||
assert params.chat_template_kwargs["enable_thinking"] is False
|
||||
|
||||
request = ChatCompletionRequest(
|
||||
model="deepseek-v4",
|
||||
messages=[{"role": "user", "content": "hi"}],
|
||||
reasoning_effort="high",
|
||||
chat_template_kwargs={"enable_thinking": False},
|
||||
)
|
||||
params = request.build_chat_params(None, "auto")
|
||||
assert params.chat_template_kwargs["enable_thinking"] is False
|
||||
|
||||
|
||||
def test_deepseek_v4_tokenizer_maps_latest_reasoning_effort_values(monkeypatch):
|
||||
captured_kwargs = []
|
||||
|
||||
def fake_encode_messages(messages, **kwargs):
|
||||
captured_kwargs.append(kwargs)
|
||||
return "prompt"
|
||||
|
||||
monkeypatch.setattr(deepseek_v4, "encode_messages", fake_encode_messages)
|
||||
tokenizer = deepseek_v4.get_deepseek_v4_tokenizer(FakeTokenizer())
|
||||
|
||||
cases = [
|
||||
("none", "chat", None),
|
||||
("minimal", "thinking", "high"),
|
||||
("low", "thinking", "high"),
|
||||
("medium", "thinking", "high"),
|
||||
("high", "thinking", "high"),
|
||||
("xhigh", "thinking", "max"),
|
||||
("max", "thinking", "max"),
|
||||
("unexpected", "thinking", "high"),
|
||||
]
|
||||
for reasoning_effort, expected_mode, expected_effort in cases:
|
||||
tokenizer.apply_chat_template(
|
||||
[{"role": "user", "content": "hi"}],
|
||||
tokenize=False,
|
||||
enable_thinking=True,
|
||||
reasoning_effort=reasoning_effort,
|
||||
)
|
||||
assert captured_kwargs[-1]["thinking_mode"] == expected_mode
|
||||
assert captured_kwargs[-1]["reasoning_effort"] == expected_effort
|
||||
140
tests/ut/patch/platform/test_patch_async_swa_kv_lifetime.py
Normal file
140
tests/ut/patch/platform/test_patch_async_swa_kv_lifetime.py
Normal file
@@ -0,0 +1,140 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
|
||||
from types import SimpleNamespace
|
||||
|
||||
import torch
|
||||
from vllm.v1.core.single_type_kv_cache_manager import MambaManager
|
||||
from vllm.v1.kv_cache_interface import MambaSpec, SlidingWindowSpec
|
||||
|
||||
import vllm_ascend.patch.platform.patch_async_swa_kv_lifetime as patch
|
||||
|
||||
|
||||
def test_schedule_output_tracks_in_flight_tokens(monkeypatch):
|
||||
request = SimpleNamespace(num_in_flight_tokens=0)
|
||||
scheduler = SimpleNamespace(requests={"request": request})
|
||||
scheduler_output = SimpleNamespace(num_scheduled_tokens={"request": 3})
|
||||
|
||||
monkeypatch.setattr(patch, "_original_update_after_schedule", lambda *_args: None)
|
||||
monkeypatch.setattr(patch, "_original_update_from_output", lambda *_args: "output")
|
||||
|
||||
patch._patched_update_after_schedule(scheduler, scheduler_output)
|
||||
assert request.num_in_flight_tokens == 3
|
||||
|
||||
assert patch._patched_update_from_output(scheduler, scheduler_output, SimpleNamespace()) == "output"
|
||||
assert request.num_in_flight_tokens == 0
|
||||
|
||||
|
||||
def test_allocate_prunes_on_processed_token_basis(monkeypatch):
|
||||
pruned_at = []
|
||||
swa_manager = SimpleNamespace(
|
||||
kv_cache_spec=SlidingWindowSpec(
|
||||
block_size=16,
|
||||
num_kv_heads=1,
|
||||
head_size=1,
|
||||
dtype=torch.float32,
|
||||
sliding_window=512,
|
||||
)
|
||||
)
|
||||
mamba_manager = MambaManager.__new__(MambaManager)
|
||||
mamba_manager.kv_cache_spec = MambaSpec(
|
||||
block_size=1,
|
||||
shapes=((1,),),
|
||||
dtypes=(torch.float32,),
|
||||
num_speculative_blocks=1,
|
||||
)
|
||||
mamba_manager.num_speculative_blocks = 1
|
||||
mamba_manager.mamba_cache_mode = "none"
|
||||
request = SimpleNamespace(
|
||||
request_id="request",
|
||||
num_in_flight_tokens=1,
|
||||
)
|
||||
|
||||
monkeypatch.setattr(
|
||||
patch,
|
||||
"_original_remove_skipped_blocks",
|
||||
lambda manager, _request_id, num_tokens: pruned_at.append((type(manager.kv_cache_spec), num_tokens)),
|
||||
)
|
||||
|
||||
def original_allocate_slots(_self, current_request):
|
||||
patch._patched_remove_skipped_blocks(swa_manager, current_request.request_id, 159)
|
||||
mamba_manager.remove_skipped_blocks(current_request.request_id, 159)
|
||||
|
||||
monkeypatch.setattr(patch, "_original_allocate_slots", original_allocate_slots)
|
||||
|
||||
patch._patched_allocate_slots(SimpleNamespace(), request)
|
||||
assert pruned_at == [(SlidingWindowSpec, 158), (MambaSpec, 158)]
|
||||
|
||||
patch._patched_remove_skipped_blocks(swa_manager, request.request_id, 159)
|
||||
assert pruned_at == [
|
||||
(SlidingWindowSpec, 158),
|
||||
(MambaSpec, 158),
|
||||
(SlidingWindowSpec, 159),
|
||||
]
|
||||
|
||||
|
||||
def test_connector_prunes_on_processed_token_basis(monkeypatch):
|
||||
pruned_at = []
|
||||
manager = SimpleNamespace(
|
||||
kv_cache_spec=SlidingWindowSpec(
|
||||
block_size=16,
|
||||
num_kv_heads=1,
|
||||
head_size=1,
|
||||
dtype=torch.float32,
|
||||
sliding_window=512,
|
||||
)
|
||||
)
|
||||
request = SimpleNamespace(
|
||||
request_id="request",
|
||||
num_in_flight_tokens=1,
|
||||
)
|
||||
|
||||
monkeypatch.setattr(
|
||||
patch,
|
||||
"_original_remove_skipped_blocks",
|
||||
lambda _self, _request_id, num_tokens: pruned_at.append(num_tokens),
|
||||
)
|
||||
|
||||
def original_connector_finished(_self, current_request):
|
||||
patch._patched_remove_skipped_blocks(manager, current_request.request_id, 159)
|
||||
return False, None
|
||||
|
||||
monkeypatch.setattr(patch, "_original_connector_finished", original_connector_finished)
|
||||
|
||||
assert patch._patched_connector_finished(SimpleNamespace(), request) == (
|
||||
False,
|
||||
None,
|
||||
)
|
||||
assert pruned_at == [158]
|
||||
|
||||
|
||||
def test_swa_admission_accounts_for_concurrent_batches(monkeypatch):
|
||||
spec = SlidingWindowSpec(
|
||||
block_size=16,
|
||||
num_kv_heads=1,
|
||||
head_size=1,
|
||||
dtype=torch.float32,
|
||||
sliding_window=512,
|
||||
)
|
||||
manager = SimpleNamespace(
|
||||
kv_cache_spec=spec,
|
||||
_max_admission_blocks_per_request=None,
|
||||
)
|
||||
vllm_config = SimpleNamespace(
|
||||
max_concurrent_batches=2,
|
||||
scheduler_config=SimpleNamespace(max_num_batched_tokens=512),
|
||||
model_config=SimpleNamespace(max_model_len=2048),
|
||||
parallel_config=SimpleNamespace(decode_context_parallel_size=1),
|
||||
)
|
||||
|
||||
def original_scheduler_init(scheduler, _vllm_config):
|
||||
scheduler.max_model_len = 2048
|
||||
scheduler.kv_cache_manager = SimpleNamespace(coordinator=SimpleNamespace(single_type_managers=(manager,)))
|
||||
|
||||
monkeypatch.setattr(patch, "_original_scheduler_init", original_scheduler_init)
|
||||
|
||||
scheduler = SimpleNamespace()
|
||||
patch._patched_scheduler_init(scheduler, vllm_config)
|
||||
|
||||
assert manager._max_admission_blocks_per_request == 97
|
||||
assert spec.max_memory_usage_bytes(vllm_config) == 97 * spec.page_size_bytes
|
||||
201
tests/ut/patch/platform/test_patch_balance_schedule.py
Normal file
201
tests/ut/patch/platform/test_patch_balance_schedule.py
Normal file
@@ -0,0 +1,201 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
|
||||
from types import SimpleNamespace
|
||||
from unittest.mock import MagicMock, patch
|
||||
|
||||
import pytest
|
||||
from vllm.v1.core.sched.interface import PauseState
|
||||
from vllm.v1.core.sched.request_queue import SchedulingPolicy
|
||||
|
||||
from vllm_ascend.patch.platform.patch_balance_schedule import (
|
||||
_ORIGINAL_SCHEDULER,
|
||||
BalanceScheduler,
|
||||
_disable_preemption_on_prefill_node,
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("kv_role", "expected"),
|
||||
[
|
||||
("kv_producer", True),
|
||||
("kv_consumer", False),
|
||||
("kv_both", False),
|
||||
(None, False),
|
||||
],
|
||||
)
|
||||
def test_disable_preemption_only_on_v023_prefill_nodes(kv_role, expected):
|
||||
kv_transfer_config = None if kv_role is None else SimpleNamespace(kv_role=kv_role)
|
||||
vllm_config = SimpleNamespace(kv_transfer_config=kv_transfer_config)
|
||||
|
||||
with patch(
|
||||
"vllm_ascend.patch.platform.patch_balance_schedule.vllm_version_is",
|
||||
return_value=True,
|
||||
):
|
||||
assert _disable_preemption_on_prefill_node(vllm_config) is expected
|
||||
|
||||
|
||||
def test_disable_preemption_is_limited_to_v023():
|
||||
vllm_config = SimpleNamespace(kv_transfer_config=SimpleNamespace(kv_role="kv_producer"))
|
||||
|
||||
with patch(
|
||||
"vllm_ascend.patch.platform.patch_balance_schedule.vllm_version_is",
|
||||
return_value=False,
|
||||
):
|
||||
assert not _disable_preemption_on_prefill_node(vllm_config)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("scheduler_name", ["default", "profiling_chunk"])
|
||||
def test_prefill_node_keeps_running_request_when_allocation_fails(
|
||||
scheduler_name,
|
||||
):
|
||||
request = SimpleNamespace(
|
||||
request_id="prefill-request",
|
||||
num_output_placeholders=0,
|
||||
num_tokens_with_spec=2,
|
||||
num_computed_tokens=1,
|
||||
num_prompt_tokens=2,
|
||||
has_encoder_inputs=False,
|
||||
spec_token_ids=[],
|
||||
)
|
||||
if scheduler_name == "default":
|
||||
scheduler = BalanceScheduler.__new__(BalanceScheduler)
|
||||
scheduler._balance_enabled = False
|
||||
else:
|
||||
from vllm_ascend.core.scheduler_profiling_chunk import (
|
||||
ProfilingChunkScheduler,
|
||||
)
|
||||
|
||||
scheduler = ProfilingChunkScheduler.__new__(ProfilingChunkScheduler)
|
||||
scheduler.profiling_chunk_manager = SimpleNamespace(
|
||||
predictor=SimpleNamespace(target_latency=None),
|
||||
is_ready=False,
|
||||
)
|
||||
scheduler.needs_kv_cache_zeroing = False
|
||||
scheduler._disable_preemption = True
|
||||
scheduler._pause_state = PauseState.UNPAUSED
|
||||
scheduler.running = [request]
|
||||
scheduler.waiting = []
|
||||
scheduler.skipped_waiting = []
|
||||
scheduler.policy = SchedulingPolicy.FCFS
|
||||
scheduler.max_num_scheduled_tokens = 1
|
||||
scheduler.max_num_encoder_input_tokens = 0
|
||||
scheduler.max_num_running_reqs = 1
|
||||
scheduler.max_model_len = 16
|
||||
scheduler.num_lookahead_tokens = 0
|
||||
scheduler.need_mamba_block_aligned_split = False
|
||||
scheduler.scheduler_config = SimpleNamespace(long_prefill_token_threshold=0)
|
||||
scheduler.kv_cache_config = SimpleNamespace(kv_cache_groups=[object()])
|
||||
scheduler.kv_cache_manager = MagicMock()
|
||||
scheduler.kv_cache_manager.allocate_slots.return_value = None
|
||||
scheduler.kv_cache_manager.get_num_common_prefix_blocks.return_value = [0]
|
||||
scheduler.encoder_cache_manager = MagicMock()
|
||||
scheduler.encoder_cache_manager.get_freed_mm_hashes.return_value = []
|
||||
scheduler.connector = None
|
||||
scheduler.ec_connector = None
|
||||
scheduler.connector_prefix_cache_stats = None
|
||||
scheduler.lora_config = None
|
||||
scheduler.is_encoder_decoder = False
|
||||
scheduler.log_stats = False
|
||||
scheduler.use_eagle = False
|
||||
scheduler.use_v2_model_runner = False
|
||||
scheduler.finished_req_ids = set()
|
||||
scheduler.prev_step_scheduled_req_ids = set()
|
||||
scheduler._preempt_request = MagicMock()
|
||||
scheduler._make_cached_request_data = MagicMock()
|
||||
scheduler._update_after_schedule = MagicMock()
|
||||
|
||||
output = scheduler.schedule()
|
||||
|
||||
scheduler.kv_cache_manager.allocate_slots.assert_called_once_with(
|
||||
request,
|
||||
1,
|
||||
num_lookahead_tokens=0,
|
||||
)
|
||||
scheduler._preempt_request.assert_not_called()
|
||||
assert scheduler.running == [request]
|
||||
assert output.total_num_scheduled_tokens == 0
|
||||
assert output.preempted_req_ids == set()
|
||||
|
||||
|
||||
def test_non_prefill_node_uses_upstream_scheduler():
|
||||
scheduler = BalanceScheduler.__new__(BalanceScheduler)
|
||||
scheduler._balance_enabled = False
|
||||
scheduler._disable_preemption = False
|
||||
expected = object()
|
||||
|
||||
with (
|
||||
patch(
|
||||
"vllm_ascend.patch.platform.patch_balance_schedule.vllm_version_is",
|
||||
return_value=True,
|
||||
),
|
||||
patch.object(_ORIGINAL_SCHEDULER, "schedule", return_value=expected) as schedule,
|
||||
):
|
||||
assert scheduler.schedule() is expected
|
||||
|
||||
schedule.assert_called_once_with()
|
||||
|
||||
|
||||
def test_async_scheduler_inherits_prefill_preemption_guard():
|
||||
from vllm.v1.core.sched.async_scheduler import AsyncScheduler
|
||||
|
||||
assert BalanceScheduler in AsyncScheduler.__mro__
|
||||
|
||||
|
||||
def test_profiling_chunk_scheduler_inherits_prefill_preemption_guard():
|
||||
from vllm_ascend.core.scheduler_profiling_chunk import (
|
||||
ProfilingChunkScheduler,
|
||||
)
|
||||
|
||||
assert BalanceScheduler in ProfilingChunkScheduler.__mro__
|
||||
|
||||
|
||||
def test_prefill_node_rejects_forced_prefix_cache_reset_while_running():
|
||||
scheduler = BalanceScheduler.__new__(BalanceScheduler)
|
||||
scheduler._disable_preemption = True
|
||||
scheduler.running = [object()]
|
||||
scheduler._preempt_request = MagicMock()
|
||||
|
||||
with pytest.raises(RuntimeError, match="drain or abort"):
|
||||
scheduler.reset_prefix_cache(reset_running_requests=True)
|
||||
|
||||
scheduler._preempt_request.assert_not_called()
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("disable_preemption", "running"),
|
||||
[
|
||||
(True, []),
|
||||
(False, [object()]),
|
||||
],
|
||||
)
|
||||
def test_forced_prefix_cache_reset_delegates_when_safe(
|
||||
disable_preemption,
|
||||
running,
|
||||
):
|
||||
scheduler = BalanceScheduler.__new__(BalanceScheduler)
|
||||
scheduler._disable_preemption = disable_preemption
|
||||
scheduler.running = running
|
||||
|
||||
with patch.object(
|
||||
_ORIGINAL_SCHEDULER,
|
||||
"reset_prefix_cache",
|
||||
return_value=True,
|
||||
) as reset_prefix_cache:
|
||||
assert scheduler.reset_prefix_cache(True, True)
|
||||
|
||||
reset_prefix_cache.assert_called_once_with(True, True)
|
||||
|
||||
|
||||
def test_non_forced_prefix_cache_reset_keeps_upstream_behavior():
|
||||
scheduler = BalanceScheduler.__new__(BalanceScheduler)
|
||||
scheduler._disable_preemption = True
|
||||
scheduler.running = [object()]
|
||||
|
||||
with patch.object(
|
||||
_ORIGINAL_SCHEDULER,
|
||||
"reset_prefix_cache",
|
||||
return_value=False,
|
||||
) as reset_prefix_cache:
|
||||
assert not scheduler.reset_prefix_cache()
|
||||
|
||||
reset_prefix_cache.assert_called_once_with(False, False)
|
||||
@@ -0,0 +1,351 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
|
||||
import json
|
||||
from unittest.mock import MagicMock
|
||||
|
||||
from vllm.entrypoints.openai.chat_completion.protocol import ChatCompletionRequest
|
||||
from vllm.tool_parsers.deepseekv4_tool_parser import DeepSeekV4ToolParser
|
||||
|
||||
from vllm_ascend.patch.platform import patch_deepseek_v4_tool_call_parser
|
||||
|
||||
MOCK_TOKENIZER = MagicMock()
|
||||
MOCK_TOKENIZER.get_vocab.return_value = {}
|
||||
|
||||
TC_START = "<｜DSML｜tool_calls>"
|
||||
TC_END = "</｜DSML｜tool_calls>"
|
||||
INV_START = '<｜DSML｜invoke name="'
|
||||
INV_END = "</｜DSML｜invoke>"
|
||||
PARAM_START = '<｜DSML｜parameter name="'
|
||||
PARAM_END = "</｜DSML｜parameter>"
|
||||
|
||||
|
||||
def _build_tool_call(
|
||||
function_name: str,
|
||||
tool_args: dict[str, str | int | bool | list[str]],
|
||||
) -> str:
|
||||
params = []
|
||||
for key, value in tool_args.items():
|
||||
if isinstance(value, bool):
|
||||
value = "false" if value is False else "true"
|
||||
string_attr = "false"
|
||||
elif isinstance(value, int):
|
||||
value = str(value)
|
||||
string_attr = "false"
|
||||
elif isinstance(value, list):
|
||||
value = json.dumps(value, ensure_ascii=False)
|
||||
string_attr = "false"
|
||||
else:
|
||||
value = str(value)
|
||||
string_attr = "true"
|
||||
|
||||
params.append(f'{PARAM_START}{key}" string="{string_attr}">{value}{PARAM_END}\n')
|
||||
|
||||
return f'{TC_START}\n{INV_START}{function_name}">\n' + "".join(params) + f"{INV_END}\n{TC_END}"
|
||||
|
||||
|
||||
def _stream(
|
||||
parser: DeepSeekV4ToolParser,
|
||||
full_text: str,
|
||||
chunk_size: int = 5,
|
||||
tools=None,
|
||||
):
|
||||
deltas = []
|
||||
previous_text = ""
|
||||
for start in range(0, len(full_text), chunk_size):
|
||||
delta_text = full_text[start : start + chunk_size]
|
||||
current_text = previous_text + delta_text
|
||||
delta = parser.extract_tool_calls_streaming(
|
||||
previous_text=previous_text,
|
||||
current_text=current_text,
|
||||
delta_text=delta_text,
|
||||
previous_token_ids=[],
|
||||
current_token_ids=[],
|
||||
delta_token_ids=[1],
|
||||
request=ChatCompletionRequest(
|
||||
model="deepseek-ai/DeepSeek-V2-Chat",
|
||||
messages=[],
|
||||
tools=tools or [_tools()],
|
||||
),
|
||||
)
|
||||
previous_text = current_text
|
||||
if delta is not None:
|
||||
deltas.append(delta)
|
||||
assert not parser._pending_delta_messages
|
||||
return deltas
|
||||
|
||||
|
||||
def _tools():
|
||||
return {
|
||||
"type": "function",
|
||||
"function": {
|
||||
"name": "plan_trip",
|
||||
"parameters": {
|
||||
"type": "object",
|
||||
"properties": {
|
||||
"days": {"type": "integer"},
|
||||
"flexible": {"type": "boolean"},
|
||||
"cities": {"type": "array", "items": {"type": "string"}},
|
||||
"notes": {"type": "string"},
|
||||
},
|
||||
"required": ["days", "flexible", "cities", "notes"],
|
||||
},
|
||||
},
|
||||
}
|
||||
|
||||
|
||||
def test_streaming_deepseek_v4_tool_calls_emit_chunked_arguments():
|
||||
parser = DeepSeekV4ToolParser(MOCK_TOKENIZER)
|
||||
full_text = _build_tool_call(
|
||||
"plan_trip",
|
||||
{
|
||||
"days": 3,
|
||||
"flexible": False,
|
||||
"cities": ["Beijing", "Shanghai", "Tokyo", "New York"],
|
||||
"notes": "靠窗座位",
|
||||
},
|
||||
)
|
||||
|
||||
deltas = _stream(parser, full_text, chunk_size=4)
|
||||
tool_chunks = []
|
||||
for delta in deltas:
|
||||
for tc in delta.tool_calls or []:
|
||||
if tc.index == 0 and tc.function and tc.function.arguments is not None:
|
||||
tool_chunks.append(tc.function.arguments)
|
||||
|
||||
reconstructed = "".join(tool_chunks)
|
||||
assert json.loads(reconstructed) == {
|
||||
"days": 3,
|
||||
"flexible": False,
|
||||
"cities": ["Beijing", "Shanghai", "Tokyo", "New York"],
|
||||
"notes": "靠窗座位",
|
||||
}
|
||||
|
||||
arg_chunks = [
|
||||
tc.function.arguments
|
||||
for delta in deltas
|
||||
for tc in delta.tool_calls or []
|
||||
if tc.index == 0 and tc.function and tc.function.arguments not in (None, "")
|
||||
]
|
||||
assert len(arg_chunks) >= 2
|
||||
|
||||
|
||||
def test_streaming_tool_call_metadata_only_first_chunk():
|
||||
parser = DeepSeekV4ToolParser(MOCK_TOKENIZER)
|
||||
full_text = _build_tool_call(
|
||||
"plan_trip",
|
||||
{
|
||||
"days": 3,
|
||||
"flexible": False,
|
||||
"cities": ["Beijing"],
|
||||
"notes": "靠窗座位",
|
||||
},
|
||||
)
|
||||
|
||||
deltas = _stream(parser, full_text, chunk_size=3)
|
||||
header_chunks = [delta for delta in deltas if delta.tool_calls]
|
||||
assert len(header_chunks) >= 1
|
||||
first = header_chunks[0].tool_calls[0]
|
||||
assert first.id is not None
|
||||
assert first.type == "function"
|
||||
assert first.function and first.function.name == "plan_trip"
|
||||
|
||||
for delta in header_chunks[1:]:
|
||||
tc = delta.tool_calls[0]
|
||||
assert tc.id is None
|
||||
if tc.function:
|
||||
assert tc.function.name is None
|
||||
assert tc.function.arguments is not None
|
||||
|
||||
|
||||
def test_streaming_wrapper_param_arguments_fragment():
|
||||
parser = DeepSeekV4ToolParser(MOCK_TOKENIZER)
|
||||
full_text = (
|
||||
TC_START
|
||||
+ "\n"
|
||||
+ f'{INV_START}plan_trip">\n'
|
||||
+ PARAM_START
|
||||
+ '__vllm_param_arguments__" string="false">{'
|
||||
+ '"days":3,"flexible":false,'
|
||||
+ '"cities":["Beijing","Shanghai","Tokyo","New York"],"notes":"靠窗座位"}</｜DSML｜parameter>\n'
|
||||
+ INV_END
|
||||
+ "\n"
|
||||
+ TC_END
|
||||
)
|
||||
|
||||
deltas = _stream(parser, full_text, chunk_size=6)
|
||||
arg_chunks = [
|
||||
tc.function.arguments
|
||||
for delta in deltas
|
||||
for tc in delta.tool_calls or []
|
||||
if tc.index == 0 and tc.function and tc.function.arguments is not None
|
||||
]
|
||||
|
||||
reconstructed = "".join(arg_chunks)
|
||||
assert json.loads(reconstructed) == {
|
||||
"days": 3,
|
||||
"flexible": False,
|
||||
"cities": ["Beijing", "Shanghai", "Tokyo", "New York"],
|
||||
"notes": "靠窗座位",
|
||||
}
|
||||
assert len(reconstructed) > 0
|
||||
|
||||
|
||||
def test_streaming_full_tool_call_single_chunk_drains_all_deltas():
|
||||
parser = DeepSeekV4ToolParser(MOCK_TOKENIZER)
|
||||
full_text = _build_tool_call(
|
||||
"plan_trip",
|
||||
{
|
||||
"days": 3,
|
||||
"flexible": False,
|
||||
"cities": ["Beijing", "Shanghai"],
|
||||
"notes": "靠窗座位",
|
||||
},
|
||||
)
|
||||
|
||||
delta = parser.extract_tool_calls_streaming(
|
||||
previous_text="",
|
||||
current_text=full_text,
|
||||
delta_text=full_text,
|
||||
previous_token_ids=[],
|
||||
current_token_ids=[],
|
||||
delta_token_ids=[1],
|
||||
request=ChatCompletionRequest(
|
||||
model="deepseek-ai/DeepSeek-V2-Chat",
|
||||
messages=[],
|
||||
tools=[_tools()],
|
||||
),
|
||||
)
|
||||
|
||||
assert delta is not None
|
||||
assert not parser._pending_delta_messages
|
||||
assert delta.tool_calls
|
||||
tool_call = delta.tool_calls[0]
|
||||
assert tool_call.id is not None
|
||||
assert tool_call.type == "function"
|
||||
assert tool_call.function and tool_call.function.name == "plan_trip"
|
||||
assert json.loads(tool_call.function.arguments) == {
|
||||
"days": 3,
|
||||
"flexible": False,
|
||||
"cities": ["Beijing", "Shanghai"],
|
||||
"notes": "靠窗座位",
|
||||
}
|
||||
|
||||
|
||||
def test_streaming_matches_non_streaming_conversion_fallbacks():
|
||||
tools = [
|
||||
{
|
||||
"type": "function",
|
||||
"function": {
|
||||
"name": "coerce",
|
||||
"parameters": {
|
||||
"type": "object",
|
||||
"properties": {
|
||||
"union_value": {"type": ["null", "string"]},
|
||||
"bad_int": {"type": "integer"},
|
||||
"nullable_string": {"type": ["null", "string"]},
|
||||
"null_string": {"type": "string"},
|
||||
"whole_number": {"type": "number"},
|
||||
},
|
||||
},
|
||||
},
|
||||
}
|
||||
]
|
||||
full_text = (
|
||||
f"{TC_START}\n"
|
||||
f'{INV_START}coerce">\n'
|
||||
f'{PARAM_START}union_value" string="false">hello{PARAM_END}\n'
|
||||
f'{PARAM_START}bad_int" string="false">abc{PARAM_END}\n'
|
||||
f'{PARAM_START}nullable_string" string="false">null{PARAM_END}\n'
|
||||
f'{PARAM_START}null_string" string="false">null{PARAM_END}\n'
|
||||
f'{PARAM_START}whole_number" string="false">3.0{PARAM_END}\n'
|
||||
f"{INV_END}\n"
|
||||
f"{TC_END}"
|
||||
)
|
||||
request = ChatCompletionRequest(
|
||||
model="deepseek-ai/DeepSeek-V2-Chat",
|
||||
messages=[],
|
||||
tools=tools,
|
||||
)
|
||||
|
||||
non_streaming = DeepSeekV4ToolParser(MOCK_TOKENIZER).extract_tool_calls(full_text, request)
|
||||
deltas = _stream(DeepSeekV4ToolParser(MOCK_TOKENIZER), full_text, chunk_size=4, tools=tools)
|
||||
|
||||
stream_args = json.loads(
|
||||
"".join(
|
||||
tc.function.arguments
|
||||
for delta in deltas
|
||||
for tc in delta.tool_calls or []
|
||||
if tc.index == 0 and tc.function and tc.function.arguments is not None
|
||||
)
|
||||
)
|
||||
expected = {
|
||||
"union_value": "hello",
|
||||
"bad_int": "abc",
|
||||
"nullable_string": None,
|
||||
"null_string": "null",
|
||||
"whole_number": 3,
|
||||
}
|
||||
assert stream_args == expected
|
||||
assert json.loads(non_streaming.tool_calls[0].function.arguments) == expected
|
||||
|
||||
|
||||
def test_composed_schema_conversion_in_streaming():
|
||||
tools = [
|
||||
{
|
||||
"type": "function",
|
||||
"function": {
|
||||
"name": "set_timer",
|
||||
"parameters": {
|
||||
"type": "object",
|
||||
"properties": {
|
||||
"wait": {
|
||||
"anyOf": [
|
||||
{"type": "object"},
|
||||
{"type": "null"},
|
||||
],
|
||||
},
|
||||
"patches": {
|
||||
"allOf": [
|
||||
{"type": "array", "items": {"type": "object"}},
|
||||
],
|
||||
},
|
||||
},
|
||||
},
|
||||
},
|
||||
}
|
||||
]
|
||||
full_text = (
|
||||
f"{TC_START}\n"
|
||||
f'{INV_START}set_timer">\n'
|
||||
f'{PARAM_START}wait" string="false">'
|
||||
f'{{"type":"for","minutes":2880}}'
|
||||
f"{PARAM_END}\n"
|
||||
f'{PARAM_START}patches" string="false">'
|
||||
f'[{{"op":"replace","path":"/schedule","value":"quiet"}}]'
|
||||
f"{PARAM_END}\n"
|
||||
f"{INV_END}\n"
|
||||
f"{TC_END}"
|
||||
)
|
||||
|
||||
deltas = _stream(DeepSeekV4ToolParser(MOCK_TOKENIZER), full_text, chunk_size=5, tools=tools)
|
||||
args = json.loads(
|
||||
"".join(
|
||||
tc.function.arguments
|
||||
for delta in deltas
|
||||
for tc in delta.tool_calls or []
|
||||
if tc.index == 0 and tc.function and tc.function.arguments is not None
|
||||
)
|
||||
)
|
||||
|
||||
assert args == {
|
||||
"wait": {"type": "for", "minutes": 2880},
|
||||
"patches": [{"op": "replace", "path": "/schedule", "value": "quiet"}],
|
||||
}
|
||||
|
||||
|
||||
def test_registered_parser_is_patch_loaded():
|
||||
# Regression check that Ascend patch applies at import-time.
|
||||
assert (
|
||||
DeepSeekV4ToolParser.extract_tool_calls_streaming
|
||||
is patch_deepseek_v4_tool_call_parser._patched_extract_tool_calls_streaming
|
||||
)
|
||||
139
tests/ut/patch/platform/test_patch_glm47_tool_call_parser.py
Normal file
139
tests/ut/patch/platform/test_patch_glm47_tool_call_parser.py
Normal file
@@ -0,0 +1,139 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
|
||||
import json
|
||||
from unittest.mock import MagicMock
|
||||
|
||||
import pytest
|
||||
|
||||
from vllm_ascend.utils import vllm_version_is
|
||||
|
||||
if not vllm_version_is("0.23.0"):
|
||||
pytest.skip(
|
||||
"upstream vLLM renamed _extract_tool_call_regions",
|
||||
allow_module_level=True,
|
||||
)
|
||||
|
||||
from vllm.entrypoints.openai.chat_completion.protocol import ChatCompletionRequest # noqa: E402
|
||||
from vllm.entrypoints.openai.chat_completion.serving import OpenAIServingChat # noqa: E402
|
||||
|
||||
# vLLM main removed the ``_WrappedParser`` helper; the base ``Parser``
|
||||
# already instantiates from ``reasoning_parser_cls`` / ``tool_parser_cls``
|
||||
# class attributes, so a thin ``DelegatingParser`` subclass is equivalent.
|
||||
from vllm.parser.abstract_parser import DelegatingParser # type: ignore[import-not-found] # noqa: E402
|
||||
from vllm.reasoning.deepseek_v3_reasoning_parser import ( # noqa: E402
|
||||
DeepSeekV3ReasoningWithThinkingParser,
|
||||
)
|
||||
from vllm.tool_parsers.glm47_moe_tool_parser import Glm47MoeModelToolParser # noqa: E402
|
||||
|
||||
from vllm_ascend.patch.platform import patch_glm47_tool_call_parser # noqa: F401, E402
|
||||
|
||||
|
||||
class _WrappedParser(DelegatingParser):
|
||||
pass
|
||||
|
||||
|
||||
MOCK_TOKENIZER = MagicMock()
|
||||
MOCK_TOKENIZER.get_vocab.return_value = {
|
||||
"<think>": 154841,
|
||||
"</think>": 154842,
|
||||
"<tool_call>": 154843,
|
||||
"</tool_call>": 154844,
|
||||
"<arg_key>": 154847,
|
||||
"</arg_key>": 154848,
|
||||
"<arg_value>": 154849,
|
||||
"</arg_value>": 154850,
|
||||
}
|
||||
|
||||
|
||||
def _request():
|
||||
return ChatCompletionRequest(
|
||||
model="glm5",
|
||||
messages=[{"role": "user", "content": "What time is it?"}],
|
||||
tools=[
|
||||
{
|
||||
"type": "function",
|
||||
"function": {
|
||||
"name": "get_current_time",
|
||||
"description": "Get the current date and time",
|
||||
"parameters": {
|
||||
"type": "object",
|
||||
"properties": {},
|
||||
},
|
||||
},
|
||||
}
|
||||
],
|
||||
tool_choice="auto",
|
||||
)
|
||||
|
||||
|
||||
def _collect_tool_args(tool_calls):
|
||||
return "".join(tc.function.arguments for tc in tool_calls if tc.function.arguments)
|
||||
|
||||
|
||||
def _parse_delta(parser, *args, finished=False, **kwargs):
|
||||
return parser.parse_delta(*args, finished=finished, **kwargs)
|
||||
|
||||
|
||||
def test_glm47_streaming_inline_zero_arg_tool_call_waits_until_complete():
|
||||
request = _request()
|
||||
parser = Glm47MoeModelToolParser(MOCK_TOKENIZER, request.tools)
|
||||
|
||||
first = parser.extract_tool_calls_streaming(
|
||||
previous_text="",
|
||||
current_text="<tool_call>get",
|
||||
delta_text="<tool_call>get",
|
||||
previous_token_ids=[],
|
||||
current_token_ids=[154843, 455],
|
||||
delta_token_ids=[154843, 455],
|
||||
request=request,
|
||||
)
|
||||
assert first is None
|
||||
|
||||
second = parser.extract_tool_calls_streaming(
|
||||
previous_text="<tool_call>get",
|
||||
current_text="<tool_call>get_current_time</tool_call>",
|
||||
delta_text="_current_time</tool_call>",
|
||||
previous_token_ids=[154843, 455],
|
||||
current_token_ids=[154843, 455, 11075, 3009, 154844],
|
||||
delta_token_ids=[11075, 3009, 154844],
|
||||
request=request,
|
||||
)
|
||||
|
||||
assert second is not None
|
||||
assert second.tool_calls
|
||||
assert second.tool_calls[0].function.name == "get_current_time"
|
||||
assert json.loads(_collect_tool_args(second.tool_calls)) == {}
|
||||
|
||||
finished = OpenAIServingChat._create_remaining_args_delta(second, "", 0)
|
||||
assert finished.tool_calls[0].function.name == "get_current_time"
|
||||
assert json.loads(_collect_tool_args(finished.tool_calls)) == {}
|
||||
|
||||
|
||||
def test_glm45_reasoning_glm47_streaming_inline_zero_arg_tool_call():
|
||||
request = _request()
|
||||
_WrappedParser.reasoning_parser_cls = DeepSeekV3ReasoningWithThinkingParser
|
||||
_WrappedParser.tool_parser_cls = Glm47MoeModelToolParser
|
||||
parser = _WrappedParser(MOCK_TOKENIZER, request.tools)
|
||||
|
||||
first = _parse_delta(
|
||||
parser,
|
||||
"Need current time.",
|
||||
[2001, 2002],
|
||||
request,
|
||||
prompt_token_ids=[],
|
||||
finished=False,
|
||||
)
|
||||
second = _parse_delta(
|
||||
parser,
|
||||
"</think><tool_call>get_current_time</tool_call>",
|
||||
[154842, 154843, 455, 11075, 3009, 154844],
|
||||
request,
|
||||
finished=True,
|
||||
)
|
||||
|
||||
assert first is not None
|
||||
assert first.reasoning == "Need current time."
|
||||
assert second is not None
|
||||
assert second.tool_calls
|
||||
assert second.tool_calls[0].function.name == "get_current_time"
|
||||
assert json.loads(_collect_tool_args(second.tool_calls)) == {}
|
||||
144
tests/ut/patch/platform/test_patch_glm_tool_call_streaming.py
Normal file
144
tests/ut/patch/platform/test_patch_glm_tool_call_streaming.py
Normal file
@@ -0,0 +1,144 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
|
||||
import json
|
||||
|
||||
from vllm.entrypoints.openai.chat_completion.serving import OpenAIServingChat
|
||||
from vllm.entrypoints.openai.engine.protocol import (
|
||||
DeltaFunctionCall,
|
||||
DeltaMessage,
|
||||
DeltaToolCall,
|
||||
)
|
||||
|
||||
from vllm_ascend.patch.platform import (
|
||||
patch_glm_tool_call_streaming as glm_streaming_patch,
|
||||
)
|
||||
|
||||
|
||||
def test_remaining_args_delta_preserves_metadata_by_default():
|
||||
original_delta = DeltaMessage(
|
||||
tool_calls=[
|
||||
DeltaToolCall(
|
||||
index=0,
|
||||
id="call_current",
|
||||
type="function",
|
||||
function=DeltaFunctionCall(
|
||||
name="current_name",
|
||||
arguments='{"files":[',
|
||||
),
|
||||
)
|
||||
]
|
||||
)
|
||||
|
||||
result = OpenAIServingChat._create_remaining_args_delta(
|
||||
original_delta,
|
||||
"]}",
|
||||
0,
|
||||
)
|
||||
|
||||
tc = result.tool_calls[0]
|
||||
assert tc.index == 0
|
||||
assert tc.id == "call_current"
|
||||
assert tc.type == "function"
|
||||
assert tc.function.name == "current_name"
|
||||
assert tc.function.arguments == "]}"
|
||||
serialized = tc.model_dump(exclude_unset=True)
|
||||
assert serialized["id"] == "call_current"
|
||||
assert serialized["type"] == "function"
|
||||
assert serialized["function"]["name"] == "current_name"
|
||||
|
||||
|
||||
def test_empty_remaining_args_delta_keeps_original_delta():
|
||||
original_delta = DeltaMessage(
|
||||
tool_calls=[
|
||||
DeltaToolCall(
|
||||
index=0,
|
||||
id="call_current",
|
||||
type="function",
|
||||
function=DeltaFunctionCall(
|
||||
name="current_name",
|
||||
arguments="",
|
||||
),
|
||||
),
|
||||
DeltaToolCall(
|
||||
index=0,
|
||||
function=DeltaFunctionCall(arguments="{}"),
|
||||
),
|
||||
]
|
||||
)
|
||||
|
||||
result = OpenAIServingChat._create_remaining_args_delta(
|
||||
original_delta,
|
||||
"",
|
||||
0,
|
||||
)
|
||||
|
||||
assert result is original_delta
|
||||
assert result.tool_calls[0].function.name == "current_name"
|
||||
assert result.tool_calls[1].function.arguments == "{}"
|
||||
|
||||
|
||||
def test_remaining_args_delta_uses_explicit_fallback_metadata():
|
||||
result = OpenAIServingChat._create_remaining_args_delta(
|
||||
DeltaMessage(),
|
||||
'{"filepath":"pong.py"}',
|
||||
0,
|
||||
fallback_tool_call_id="call_files",
|
||||
fallback_tool_call_type="function",
|
||||
fallback_tool_call_name="builtin_read_many_files",
|
||||
)
|
||||
|
||||
tc = result.tool_calls[0]
|
||||
assert tc.index == 0
|
||||
assert tc.id == "call_files"
|
||||
assert tc.type == "function"
|
||||
assert tc.function.name == "builtin_read_many_files"
|
||||
assert tc.function.arguments == '{"filepath":"pong.py"}'
|
||||
|
||||
|
||||
def test_terminal_argument_chunk_is_split_before_finish_chunk():
|
||||
chunk = {
|
||||
"id": "chatcmpl-test",
|
||||
"object": "chat.completion.chunk",
|
||||
"created": 0,
|
||||
"model": "GLM-5",
|
||||
"choices": [
|
||||
{
|
||||
"index": 0,
|
||||
"delta": {
|
||||
"tool_calls": [
|
||||
{
|
||||
"index": 0,
|
||||
"function": {
|
||||
"arguments": '"pong.py"}',
|
||||
},
|
||||
}
|
||||
]
|
||||
},
|
||||
"finish_reason": "tool_calls",
|
||||
"stop_reason": None,
|
||||
}
|
||||
],
|
||||
}
|
||||
|
||||
chunks = glm_streaming_patch._split_terminal_tool_arg_chunk(f"data: {json.dumps(chunk)}\n\n")
|
||||
|
||||
assert len(chunks) == 2
|
||||
arg_payload = json.loads(chunks[0].removeprefix("data: ").removesuffix("\n\n"))
|
||||
finish_payload = json.loads(chunks[1].removeprefix("data: ").removesuffix("\n\n"))
|
||||
|
||||
arg_choice = arg_payload["choices"][0]
|
||||
assert arg_choice["finish_reason"] is None
|
||||
assert arg_choice["stop_reason"] is None
|
||||
assert arg_choice["delta"]["tool_calls"][0]["function"]["arguments"] == '"pong.py"}'
|
||||
|
||||
finish_choice = finish_payload["choices"][0]
|
||||
assert finish_choice["finish_reason"] == "tool_calls"
|
||||
assert finish_choice["delta"] == {}
|
||||
|
||||
|
||||
def test_non_terminal_and_done_chunks_are_not_split():
|
||||
content = 'data: {"choices":[]}\n\n'
|
||||
done = "data: [DONE]\n\n"
|
||||
|
||||
assert glm_streaming_patch._split_terminal_tool_arg_chunk(content) == [content]
|
||||
assert glm_streaming_patch._split_terminal_tool_arg_chunk(done) == [done]
|
||||
@@ -0,0 +1,400 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
|
||||
import json
|
||||
from typing import Any
|
||||
|
||||
import pytest
|
||||
|
||||
from vllm_ascend.utils import vllm_version_is
|
||||
|
||||
pytestmark = pytest.mark.skipif(
|
||||
not vllm_version_is("0.23.0"),
|
||||
reason="upstream vLLM removed tool_call_start_token attribute",
|
||||
)
|
||||
|
||||
from openai.types.responses.function_tool import FunctionTool # noqa: E402
|
||||
from vllm.entrypoints.openai.chat_completion.protocol import ( # noqa: E402
|
||||
ChatCompletionToolsParam,
|
||||
FunctionDefinition,
|
||||
)
|
||||
from vllm.tool_parsers.minimax_m2_tool_parser import MinimaxM2ToolParser # noqa: E402
|
||||
|
||||
from vllm_ascend.patch.platform import ( # noqa: E402
|
||||
patch_minimax_m2_tool_call_parser as minimax_m2_patch,
|
||||
)
|
||||
|
||||
TC_START_ID = 1
|
||||
TC_END_ID = 2
|
||||
EOS_ID = 99
|
||||
|
||||
|
||||
class FakeTokenizer:
|
||||
def get_vocab(self):
|
||||
return {
|
||||
"<minimax:tool_call>": TC_START_ID,
|
||||
"</minimax:tool_call>": TC_END_ID,
|
||||
}
|
||||
|
||||
|
||||
def _feed(parser: MinimaxM2ToolParser, chunks):
|
||||
previous = ""
|
||||
results = []
|
||||
for chunk in chunks:
|
||||
if isinstance(chunk, tuple):
|
||||
delta, delta_ids = chunk
|
||||
else:
|
||||
delta = chunk
|
||||
delta_ids = []
|
||||
|
||||
current = previous + delta
|
||||
result = parser.extract_tool_calls_streaming(
|
||||
previous_text=previous,
|
||||
current_text=current,
|
||||
delta_text=delta,
|
||||
previous_token_ids=[],
|
||||
current_token_ids=[],
|
||||
delta_token_ids=delta_ids,
|
||||
request=None,
|
||||
)
|
||||
if result is not None:
|
||||
results.append(result)
|
||||
previous = current
|
||||
return results
|
||||
|
||||
|
||||
def _collect_content(results):
|
||||
return "".join(result.content for result in results if result.content)
|
||||
|
||||
|
||||
def _collect_tool_calls(results):
|
||||
tool_calls: dict[int, dict[str, Any]] = {}
|
||||
for result in results:
|
||||
for tool_call in result.tool_calls or []:
|
||||
tool_calls.setdefault(
|
||||
tool_call.index,
|
||||
{
|
||||
"id": None,
|
||||
"name": "",
|
||||
"arguments": "",
|
||||
},
|
||||
)
|
||||
if tool_call.id:
|
||||
tool_calls[tool_call.index]["id"] = tool_call.id
|
||||
if tool_call.function:
|
||||
if tool_call.function.name:
|
||||
tool_calls[tool_call.index]["name"] += tool_call.function.name
|
||||
if tool_call.function.arguments:
|
||||
tool_calls[tool_call.index]["arguments"] += tool_call.function.arguments
|
||||
return tool_calls
|
||||
|
||||
|
||||
def test_registered_parser_is_patch_loaded():
|
||||
assert MinimaxM2ToolParser.extract_tool_calls_streaming is minimax_m2_patch._patched_extract_tool_calls_streaming
|
||||
|
||||
|
||||
def test_plain_content_before_tool_call_is_preserved():
|
||||
parser = MinimaxM2ToolParser(FakeTokenizer())
|
||||
results = _feed(
|
||||
parser,
|
||||
[
|
||||
"Let me check. ",
|
||||
'<minimax:tool_call><invoke name="get_weather">'
|
||||
'<parameter name="city">Seattle</parameter>'
|
||||
"</invoke></minimax:tool_call>",
|
||||
],
|
||||
)
|
||||
|
||||
assert _collect_content(results) == "Let me check. "
|
||||
assert len(parser.prev_tool_call_arr) == 1
|
||||
|
||||
|
||||
def test_plain_content_before_partial_tool_call_omits_tool_calls_payload():
|
||||
parser = MinimaxM2ToolParser(FakeTokenizer())
|
||||
results = _feed(parser, ["Let me check. <minimax:tool_call>"])
|
||||
|
||||
assert len(results) == 1
|
||||
assert results[0].content == "Let me check. "
|
||||
assert results[0].tool_calls == []
|
||||
assert "tool_calls" not in results[0].model_dump(exclude_unset=True)
|
||||
assert "tool_calls" not in results[0].model_dump_json(exclude_unset=True)
|
||||
|
||||
|
||||
def test_streaming_emits_tool_name_before_argument_fragments():
|
||||
parser = MinimaxM2ToolParser(FakeTokenizer())
|
||||
results = _feed(
|
||||
parser,
|
||||
[
|
||||
"Let me check. ",
|
||||
"<minimax:tool_call>",
|
||||
'<invoke name="get_weather">',
|
||||
'<parameter name="city">Sea',
|
||||
"ttle</parameter>",
|
||||
"</invoke></minimax:tool_call>",
|
||||
],
|
||||
)
|
||||
|
||||
tool_deltas = [tc for result in results for tc in (result.tool_calls or [])]
|
||||
argument_fragments = [tc.function.arguments for tc in tool_deltas[1:] if tc.function and tc.function.arguments]
|
||||
|
||||
assert _collect_content(results) == "Let me check. "
|
||||
assert tool_deltas[0].function.name == "get_weather"
|
||||
assert tool_deltas[0].function.arguments is None
|
||||
assert argument_fragments == ['{"city":"Seattle"', "}"]
|
||||
assert "".join(argument_fragments) == '{"city":"Seattle"}'
|
||||
|
||||
|
||||
def test_streaming_waits_for_parameter_close_before_arguments():
|
||||
parser = MinimaxM2ToolParser(FakeTokenizer())
|
||||
results = _feed(
|
||||
parser,
|
||||
[
|
||||
"<minimax:tool_call>",
|
||||
'<invoke name="get_weather">',
|
||||
'<parameter name="city">Sea',
|
||||
],
|
||||
)
|
||||
|
||||
tool_deltas = [tc for result in results for tc in (result.tool_calls or [])]
|
||||
|
||||
assert tool_deltas[0].function.name == "get_weather"
|
||||
assert tool_deltas[0].function.arguments is None
|
||||
assert len(tool_deltas) == 1
|
||||
assert parser.prev_tool_call_arr == []
|
||||
|
||||
|
||||
def test_parameter_end_tag_token_pieces_not_streamed_as_arguments():
|
||||
parser = MinimaxM2ToolParser(
|
||||
FakeTokenizer(),
|
||||
tools=[
|
||||
ChatCompletionToolsParam(
|
||||
function=FunctionDefinition(
|
||||
name="write_file",
|
||||
parameters={
|
||||
"type": "object",
|
||||
"properties": {
|
||||
"content": {"type": "string"},
|
||||
"path": {"type": "string"},
|
||||
},
|
||||
},
|
||||
),
|
||||
)
|
||||
],
|
||||
)
|
||||
|
||||
results = _feed(
|
||||
parser,
|
||||
[
|
||||
("", [TC_START_ID]),
|
||||
"\n",
|
||||
"<",
|
||||
"invoke",
|
||||
" name",
|
||||
'="',
|
||||
"write",
|
||||
"_file",
|
||||
'">\n',
|
||||
"<",
|
||||
"parameter",
|
||||
" name",
|
||||
'="',
|
||||
"path",
|
||||
'">',
|
||||
"a",
|
||||
".txt",
|
||||
"</",
|
||||
"parameter",
|
||||
">\n",
|
||||
"<",
|
||||
"parameter",
|
||||
" name",
|
||||
'="',
|
||||
"content",
|
||||
'">',
|
||||
"123",
|
||||
"</",
|
||||
"parameter",
|
||||
">\n",
|
||||
"</",
|
||||
"invoke",
|
||||
">\n",
|
||||
("", [TC_END_ID]),
|
||||
("", [EOS_ID]),
|
||||
],
|
||||
)
|
||||
|
||||
args = _collect_tool_calls(results)[0]["arguments"]
|
||||
|
||||
assert "</parameter" not in args
|
||||
assert json.loads(args) == {"path": "a.txt", "content": "123"}
|
||||
assert parser.prev_tool_call_arr == [
|
||||
{
|
||||
"name": "write_file",
|
||||
"arguments": {"path": "a.txt", "content": "123"},
|
||||
}
|
||||
]
|
||||
|
||||
|
||||
def test_complete_single_chunk_still_reconstructs_tool_call():
|
||||
parser = MinimaxM2ToolParser(FakeTokenizer())
|
||||
results = _feed(
|
||||
parser,
|
||||
[
|
||||
'<minimax:tool_call><invoke name="get_weather">'
|
||||
'<parameter name="city">Seattle</parameter>'
|
||||
"</invoke></minimax:tool_call>",
|
||||
("", [EOS_ID]),
|
||||
],
|
||||
)
|
||||
|
||||
tool_calls = _collect_tool_calls(results)
|
||||
|
||||
assert len(tool_calls) == 1
|
||||
assert tool_calls[0]["name"] == "get_weather"
|
||||
assert json.loads(tool_calls[0]["arguments"]) == {"city": "Seattle"}
|
||||
assert results[-1].content == ""
|
||||
|
||||
|
||||
def test_start_token_can_arrive_as_special_token_id():
|
||||
parser = MinimaxM2ToolParser(FakeTokenizer())
|
||||
results = _feed(
|
||||
parser,
|
||||
[
|
||||
("", [TC_START_ID]),
|
||||
'<invoke name="get_weather">',
|
||||
'<parameter name="city">Seattle</parameter>',
|
||||
"</invoke>",
|
||||
("", [TC_END_ID]),
|
||||
("", [EOS_ID]),
|
||||
],
|
||||
)
|
||||
|
||||
tool_calls = _collect_tool_calls(results)
|
||||
|
||||
assert len(tool_calls) == 1
|
||||
assert tool_calls[0]["name"] == "get_weather"
|
||||
assert json.loads(tool_calls[0]["arguments"]) == {"city": "Seattle"}
|
||||
assert results[-1].content == ""
|
||||
|
||||
|
||||
def test_start_token_id_survives_empty_chunks_before_invoke_text():
|
||||
parser = MinimaxM2ToolParser(FakeTokenizer())
|
||||
results = _feed(
|
||||
parser,
|
||||
[
|
||||
("", [TC_START_ID]),
|
||||
("", []),
|
||||
("", []),
|
||||
'<invoke name="get_weather">',
|
||||
'<parameter name="city">Seattle</parameter>',
|
||||
"</invoke>",
|
||||
("", [TC_END_ID]),
|
||||
("", [EOS_ID]),
|
||||
],
|
||||
)
|
||||
|
||||
tool_calls = _collect_tool_calls(results)
|
||||
|
||||
assert len(tool_calls) == 1
|
||||
assert tool_calls[0]["name"] == "get_weather"
|
||||
assert json.loads(tool_calls[0]["arguments"]) == {"city": "Seattle"}
|
||||
assert results[-1].content == ""
|
||||
|
||||
|
||||
def test_chat_tool_schema_drives_type_conversion():
|
||||
parser = MinimaxM2ToolParser(
|
||||
FakeTokenizer(),
|
||||
tools=[
|
||||
ChatCompletionToolsParam(
|
||||
function=FunctionDefinition(
|
||||
name="get_weather",
|
||||
parameters={
|
||||
"type": "object",
|
||||
"properties": {"days": {"type": "integer"}},
|
||||
},
|
||||
),
|
||||
)
|
||||
],
|
||||
)
|
||||
results = _feed(
|
||||
parser,
|
||||
[
|
||||
'<minimax:tool_call><invoke name="get_weather">'
|
||||
'<parameter name="days">5</parameter>'
|
||||
"</invoke></minimax:tool_call>",
|
||||
],
|
||||
)
|
||||
|
||||
parsed = json.loads(_collect_tool_calls(results)[0]["arguments"])
|
||||
|
||||
assert parsed["days"] == 5
|
||||
assert isinstance(parsed["days"], int)
|
||||
|
||||
|
||||
def test_patch_does_not_require_private_v0202_schema_helpers(monkeypatch):
|
||||
monkeypatch.delattr(
|
||||
MinimaxM2ToolParser,
|
||||
"_get_param_types_from_config",
|
||||
raising=False,
|
||||
)
|
||||
monkeypatch.delattr(
|
||||
MinimaxM2ToolParser,
|
||||
"_convert_param_value_with_types",
|
||||
raising=False,
|
||||
)
|
||||
parser = MinimaxM2ToolParser(
|
||||
FakeTokenizer(),
|
||||
tools=[
|
||||
ChatCompletionToolsParam(
|
||||
function=FunctionDefinition(
|
||||
name="get_weather",
|
||||
parameters={
|
||||
"type": "object",
|
||||
"properties": {"days": {"type": "integer"}},
|
||||
},
|
||||
),
|
||||
)
|
||||
],
|
||||
)
|
||||
results = _feed(
|
||||
parser,
|
||||
[
|
||||
'<minimax:tool_call><invoke name="get_weather">'
|
||||
'<parameter name="days">5</parameter>'
|
||||
"</invoke></minimax:tool_call>",
|
||||
],
|
||||
)
|
||||
|
||||
parsed = json.loads(_collect_tool_calls(results)[0]["arguments"])
|
||||
|
||||
assert parsed["days"] == 5
|
||||
assert isinstance(parsed["days"], int)
|
||||
|
||||
|
||||
def test_responses_function_tool_schema_drives_type_conversion():
|
||||
parser = MinimaxM2ToolParser(
|
||||
FakeTokenizer(),
|
||||
tools=[
|
||||
FunctionTool(
|
||||
type="function",
|
||||
name="get_weather",
|
||||
description="Get weather data",
|
||||
parameters={
|
||||
"type": "object",
|
||||
"properties": {"days": {"type": "integer"}},
|
||||
},
|
||||
)
|
||||
],
|
||||
)
|
||||
results = _feed(
|
||||
parser,
|
||||
[
|
||||
'<minimax:tool_call><invoke name="get_weather">'
|
||||
'<parameter name="days">5</parameter>'
|
||||
"</invoke></minimax:tool_call>",
|
||||
],
|
||||
)
|
||||
|
||||
parsed = json.loads(_collect_tool_calls(results)[0]["arguments"])
|
||||
|
||||
assert parsed["days"] == 5
|
||||
assert isinstance(parsed["days"], int)
|
||||
414
tests/ut/patch/platform/test_patch_minimax_usage_accounting.py
Normal file
414
tests/ut/patch/platform/test_patch_minimax_usage_accounting.py
Normal file
@@ -0,0 +1,414 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
|
||||
import json
|
||||
from types import SimpleNamespace
|
||||
|
||||
import pytest
|
||||
|
||||
from vllm_ascend.utils import vllm_version_is
|
||||
|
||||
pytestmark = pytest.mark.skipif(
|
||||
not vllm_version_is("0.23.0"),
|
||||
reason="upstream vLLM removed end_token_id attribute",
|
||||
)
|
||||
from vllm.entrypoints.openai.chat_completion.serving import OpenAIServingChat # noqa: E402
|
||||
from vllm.parser.parser_manager import ParserManager # noqa: E402
|
||||
from vllm.reasoning.minimax_m2_reasoning_parser import ( # noqa: E402
|
||||
MiniMaxM2AppendThinkReasoningParser,
|
||||
MiniMaxM2ReasoningParser,
|
||||
)
|
||||
|
||||
from vllm_ascend.patch.platform import patch_minimax_usage_accounting as usage_patch # noqa: E402
|
||||
|
||||
|
||||
class FakeTokenizer:
|
||||
def get_vocab(self):
|
||||
return {
|
||||
"<think>": 1,
|
||||
"</think>": 2,
|
||||
"<minimax:tool_call>": 3,
|
||||
"</minimax:tool_call>": 4,
|
||||
}
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("parser_cls", "token_ids", "expected_reasoning_tokens"),
|
||||
[
|
||||
pytest.param(
|
||||
MiniMaxM2ReasoningParser,
|
||||
[10, 11, 2, 20],
|
||||
2,
|
||||
id="minimax-reasoning-before-end-token",
|
||||
),
|
||||
pytest.param(
|
||||
MiniMaxM2AppendThinkReasoningParser,
|
||||
[10, 11, 2, 20],
|
||||
2,
|
||||
id="append-think-reasoning-before-end-token",
|
||||
),
|
||||
pytest.param(
|
||||
MiniMaxM2ReasoningParser,
|
||||
[10, 11, 20],
|
||||
3,
|
||||
id="minimax-no-end-token-means-all-output-is-reasoning",
|
||||
),
|
||||
pytest.param(
|
||||
MiniMaxM2AppendThinkReasoningParser,
|
||||
[10, 11, 20],
|
||||
3,
|
||||
id="append-think-no-end-token-means-all-output-is-reasoning",
|
||||
),
|
||||
pytest.param(
|
||||
MiniMaxM2ReasoningParser,
|
||||
[2, 20],
|
||||
0,
|
||||
id="minimax-end-token-first-means-no-reasoning-tokens",
|
||||
),
|
||||
pytest.param(
|
||||
MiniMaxM2AppendThinkReasoningParser,
|
||||
[2, 20],
|
||||
0,
|
||||
id="append-think-end-token-first-means-no-reasoning-tokens",
|
||||
),
|
||||
],
|
||||
)
|
||||
def test_count_reasoning_tokens(
|
||||
parser_cls,
|
||||
token_ids,
|
||||
expected_reasoning_tokens,
|
||||
):
|
||||
parser = parser_cls(FakeTokenizer())
|
||||
|
||||
assert parser.count_reasoning_tokens(token_ids) == expected_reasoning_tokens
|
||||
|
||||
|
||||
def test_update_usage_tracking_state_tracks_prompt_and_completion_tokens():
|
||||
state = usage_patch._create_usage_tracking_state(
|
||||
num_choices=2,
|
||||
reasoning_parser=None,
|
||||
)
|
||||
|
||||
res = SimpleNamespace(
|
||||
prompt_token_ids=[1, 2],
|
||||
encoder_prompt_token_ids=[3],
|
||||
num_cached_tokens=4,
|
||||
outputs=[
|
||||
SimpleNamespace(index=0, token_ids=(10, 11)),
|
||||
SimpleNamespace(index=1, token_ids=[20]),
|
||||
],
|
||||
)
|
||||
|
||||
usage_patch._update_usage_tracking_state(state, res)
|
||||
|
||||
assert state.num_prompt_tokens == 3
|
||||
assert state.num_cached_tokens == 4
|
||||
assert state.completion_tokens == [2, 1]
|
||||
assert state.raw_output_token_ids == [[10, 11], [20]]
|
||||
|
||||
|
||||
def test_make_usage_info_injects_reasoning_token_details():
|
||||
fake_serving = SimpleNamespace(enable_prompt_tokens_details=True)
|
||||
usage = usage_patch._make_usage_info(
|
||||
fake_serving,
|
||||
prompt_tokens=3,
|
||||
completion_tokens=4,
|
||||
num_cached_tokens=1,
|
||||
reasoning_tokens=2,
|
||||
)
|
||||
|
||||
payload = usage.model_dump(exclude_none=True)
|
||||
|
||||
assert payload["completion_tokens_details"]["reasoning_tokens"] == 2
|
||||
assert payload["prompt_tokens_details"]["cached_tokens"] == 1
|
||||
|
||||
|
||||
def test_make_usage_info_injects_zero_cached_tokens():
|
||||
fake_serving = SimpleNamespace(enable_prompt_tokens_details=True)
|
||||
usage = usage_patch._make_usage_info(
|
||||
fake_serving,
|
||||
prompt_tokens=3,
|
||||
completion_tokens=4,
|
||||
num_cached_tokens=0,
|
||||
)
|
||||
|
||||
payload = usage.model_dump(exclude_none=True)
|
||||
|
||||
assert payload["prompt_tokens_details"]["cached_tokens"] == 0
|
||||
|
||||
|
||||
def test_make_full_response_usage_sums_reasoning_tokens():
|
||||
class FakeServing:
|
||||
enable_prompt_tokens_details = False
|
||||
|
||||
def _make_usage_info(self, **kwargs):
|
||||
return usage_patch._make_usage_info(self, **kwargs)
|
||||
|
||||
state = usage_patch._create_usage_tracking_state(
|
||||
num_choices=2,
|
||||
reasoning_parser=MiniMaxM2ReasoningParser(FakeTokenizer()),
|
||||
)
|
||||
state.num_prompt_tokens = 3
|
||||
state.num_cached_tokens = 1
|
||||
state.final_res = SimpleNamespace(num_cached_tokens=1)
|
||||
state.completion_tokens = [4, 2]
|
||||
state.raw_output_token_ids = [[10, 11, 2, 20], [30, 31]]
|
||||
|
||||
usage = usage_patch._make_full_response_usage(FakeServing(), state)
|
||||
|
||||
assert usage.prompt_tokens == 3
|
||||
assert usage.completion_tokens == 6
|
||||
assert usage.total_tokens == 9
|
||||
assert usage.completion_tokens_details.reasoning_tokens == 4
|
||||
assert usage.prompt_tokens_details is None
|
||||
|
||||
|
||||
def test_make_full_response_usage_accepts_wrapped_reasoning_parser():
|
||||
class FakeServing:
|
||||
enable_prompt_tokens_details = False
|
||||
|
||||
def _make_usage_info(self, **kwargs):
|
||||
return usage_patch._make_usage_info(self, **kwargs)
|
||||
|
||||
state = usage_patch._create_usage_tracking_state(
|
||||
num_choices=1,
|
||||
reasoning_parser=SimpleNamespace(
|
||||
reasoning_parser=MiniMaxM2ReasoningParser(FakeTokenizer()),
|
||||
),
|
||||
)
|
||||
state.num_prompt_tokens = 3
|
||||
state.final_res = SimpleNamespace(num_cached_tokens=None)
|
||||
state.completion_tokens = [4]
|
||||
state.raw_output_token_ids = [[10, 11, 2, 20]]
|
||||
|
||||
usage = usage_patch._make_full_response_usage(FakeServing(), state)
|
||||
|
||||
assert usage.completion_tokens_details.reasoning_tokens == 2
|
||||
|
||||
|
||||
def test_count_reasoning_tokens_accepts_minimax_unified_parser():
|
||||
parser_cls = ParserManager.get_parser(
|
||||
tool_parser_name="minimax_m2",
|
||||
reasoning_parser_name="minimax_m2",
|
||||
enable_auto_tools=True,
|
||||
model_name="MiniMax-M2",
|
||||
)
|
||||
parser = parser_cls(FakeTokenizer(), tools=[])
|
||||
|
||||
assert not hasattr(parser, "count_reasoning_tokens")
|
||||
assert usage_patch._count_minimax_reasoning_tokens_for_usage([10, 11, 2, 20], parser) == 2
|
||||
|
||||
|
||||
def test_count_reasoning_tokens_accepts_wrapped_minimax_parser():
|
||||
parser = SimpleNamespace(
|
||||
reasoning_parser=MiniMaxM2ReasoningParser(FakeTokenizer()),
|
||||
)
|
||||
|
||||
assert usage_patch._count_minimax_reasoning_tokens_for_usage([10, 11, 2, 20], parser) == 2
|
||||
assert usage_patch._is_minimax_reasoning_parser(parser)
|
||||
|
||||
|
||||
def test_count_reasoning_tokens_skips_non_minimax_parser_manager_wrapper():
|
||||
parser_cls = ParserManager.get_parser(
|
||||
tool_parser_name="deepseek_v4",
|
||||
reasoning_parser_name="deepseek_v4",
|
||||
enable_auto_tools=True,
|
||||
model_name="DeepSeek-V4",
|
||||
)
|
||||
parser = parser_cls(FakeTokenizer(), tools=[])
|
||||
|
||||
assert not hasattr(parser, "count_reasoning_tokens")
|
||||
assert usage_patch._count_minimax_reasoning_tokens_for_usage([10, 11], parser) is None
|
||||
assert not usage_patch._is_minimax_reasoning_parser(parser)
|
||||
|
||||
|
||||
def test_non_minimax_parser_does_not_enable_tracking_by_default():
|
||||
class FakeReasoningParser:
|
||||
def count_reasoning_tokens(self, token_ids):
|
||||
return len(token_ids)
|
||||
|
||||
parser = FakeReasoningParser()
|
||||
|
||||
assert usage_patch._count_minimax_reasoning_tokens_for_usage([10, 11], parser) is None
|
||||
assert not usage_patch._is_minimax_reasoning_parser(parser)
|
||||
assert usage_patch._sum_reasoning_tokens_for_usage([[10, 11]], parser) is None
|
||||
|
||||
|
||||
def test_make_full_response_usage_skips_non_minimax_reasoning_details():
|
||||
class FakeServing:
|
||||
enable_prompt_tokens_details = True
|
||||
|
||||
def _make_usage_info(self, **kwargs):
|
||||
return usage_patch._make_usage_info(self, **kwargs)
|
||||
|
||||
class FakeReasoningParser:
|
||||
def count_reasoning_tokens(self, token_ids):
|
||||
return len(token_ids)
|
||||
|
||||
state = usage_patch._create_usage_tracking_state(
|
||||
num_choices=1,
|
||||
reasoning_parser=FakeReasoningParser(),
|
||||
enable_prompt_tokens_details=True,
|
||||
)
|
||||
state.num_prompt_tokens = 3
|
||||
state.num_cached_tokens = 0
|
||||
state.final_res = SimpleNamespace(num_cached_tokens=0)
|
||||
state.completion_tokens = [2]
|
||||
state.raw_output_token_ids = [[10, 11]]
|
||||
|
||||
usage = usage_patch._make_full_response_usage(FakeServing(), state)
|
||||
|
||||
assert usage.completion_tokens_details is None
|
||||
assert usage.prompt_tokens_details.cached_tokens == 0
|
||||
|
||||
|
||||
def test_chat_generators_are_not_patched_at_class_level():
|
||||
assert (
|
||||
OpenAIServingChat.chat_completion_stream_generator is not usage_patch._wrapped_chat_completion_stream_generator
|
||||
)
|
||||
assert OpenAIServingChat.chat_completion_full_generator is not usage_patch._wrapped_chat_completion_full_generator
|
||||
|
||||
|
||||
def test_chat_init_is_not_wrapped_by_minimax_usage_patch():
|
||||
assert not hasattr(OpenAIServingChat, "_ascend_original_init_for_minimax_usage")
|
||||
assert "patch_minimax_usage_accounting.py" not in OpenAIServingChat.__init__.__code__.co_filename
|
||||
|
||||
|
||||
def test_reasoning_parser_cls_descriptor_preserves_default_access():
|
||||
descriptor = OpenAIServingChat.__dict__["reasoning_parser_cls"]
|
||||
serving = object.__new__(OpenAIServingChat)
|
||||
|
||||
assert OpenAIServingChat.reasoning_parser_cls is descriptor.default_value
|
||||
assert serving.reasoning_parser_cls is descriptor.default_value
|
||||
|
||||
|
||||
def test_chat_usage_wrapper_is_bound_only_for_target_instances():
|
||||
class FakeReasoningParser:
|
||||
pass
|
||||
|
||||
non_minimax_serving = SimpleNamespace(
|
||||
enable_prompt_tokens_details=False,
|
||||
reasoning_parser_cls=FakeReasoningParser,
|
||||
)
|
||||
minimax_serving = SimpleNamespace(
|
||||
enable_prompt_tokens_details=False,
|
||||
reasoning_parser_cls=MiniMaxM2ReasoningParser,
|
||||
)
|
||||
non_minimax_prompt_details_serving = SimpleNamespace(
|
||||
enable_prompt_tokens_details=True,
|
||||
reasoning_parser_cls=FakeReasoningParser,
|
||||
)
|
||||
|
||||
assert not usage_patch._should_patch_chat_usage_instance(non_minimax_serving)
|
||||
assert usage_patch._should_patch_chat_usage_instance(minimax_serving)
|
||||
assert not usage_patch._should_patch_chat_usage_instance(non_minimax_prompt_details_serving)
|
||||
|
||||
|
||||
def test_reasoning_parser_cls_assignment_binds_only_minimax_instances():
|
||||
class FakeReasoningParser:
|
||||
pass
|
||||
|
||||
non_minimax_serving = object.__new__(OpenAIServingChat)
|
||||
non_minimax_serving.reasoning_parser_cls = FakeReasoningParser
|
||||
|
||||
assert non_minimax_serving.reasoning_parser_cls is FakeReasoningParser
|
||||
assert "chat_completion_stream_generator" not in non_minimax_serving.__dict__
|
||||
assert "chat_completion_full_generator" not in non_minimax_serving.__dict__
|
||||
|
||||
minimax_serving = object.__new__(OpenAIServingChat)
|
||||
minimax_serving.reasoning_parser_cls = MiniMaxM2ReasoningParser
|
||||
|
||||
assert minimax_serving.reasoning_parser_cls is MiniMaxM2ReasoningParser
|
||||
assert (
|
||||
minimax_serving.chat_completion_stream_generator.__func__
|
||||
is usage_patch._wrapped_chat_completion_stream_generator
|
||||
)
|
||||
assert (
|
||||
minimax_serving.chat_completion_full_generator.__func__ is usage_patch._wrapped_chat_completion_full_generator
|
||||
)
|
||||
|
||||
|
||||
def test_instance_wrapper_composes_with_class_level_stream_patches():
|
||||
serving = SimpleNamespace(
|
||||
enable_prompt_tokens_details=False,
|
||||
reasoning_parser_cls=MiniMaxM2ReasoningParser,
|
||||
)
|
||||
|
||||
usage_patch._patch_chat_usage_instance(serving)
|
||||
|
||||
assert (
|
||||
serving._ascend_original_chat_completion_stream_generator.__func__
|
||||
is OpenAIServingChat.chat_completion_stream_generator
|
||||
)
|
||||
assert (
|
||||
serving._ascend_original_chat_completion_full_generator.__func__
|
||||
is OpenAIServingChat.chat_completion_full_generator
|
||||
)
|
||||
assert serving.chat_completion_stream_generator.__func__ is usage_patch._wrapped_chat_completion_stream_generator
|
||||
assert serving.chat_completion_full_generator.__func__ is usage_patch._wrapped_chat_completion_full_generator
|
||||
|
||||
|
||||
def test_stream_usage_details_are_injected_without_replacing_source():
|
||||
state = usage_patch._create_usage_tracking_state(
|
||||
num_choices=1,
|
||||
reasoning_parser=MiniMaxM2ReasoningParser(FakeTokenizer()),
|
||||
enable_prompt_tokens_details=True,
|
||||
)
|
||||
state.num_cached_tokens = 0
|
||||
state.raw_output_token_ids = [[10, 11, 2, 20]]
|
||||
|
||||
chunk = {
|
||||
"id": "chatcmpl-test",
|
||||
"object": "chat.completion.chunk",
|
||||
"choices": [{"index": 0, "delta": {}, "finish_reason": None}],
|
||||
"usage": {
|
||||
"prompt_tokens": 3,
|
||||
"completion_tokens": 4,
|
||||
"total_tokens": 7,
|
||||
},
|
||||
}
|
||||
|
||||
data = usage_patch._inject_stream_usage_details(
|
||||
f"data: {json.dumps(chunk)}\n\n",
|
||||
state,
|
||||
)
|
||||
payload = json.loads(data.removeprefix("data: ").removesuffix("\n\n"))
|
||||
|
||||
assert payload["usage"]["completion_tokens_details"] == {
|
||||
"reasoning_tokens": 2,
|
||||
}
|
||||
assert payload["usage"]["prompt_tokens_details"] == {
|
||||
"cached_tokens": 0,
|
||||
}
|
||||
assert not hasattr(usage_patch, "_extract_class_method_source")
|
||||
assert not hasattr(usage_patch, "_patch_chat_completion_stream_generator")
|
||||
|
||||
|
||||
def test_stream_usage_details_inject_prompt_details_without_reasoning():
|
||||
state = usage_patch._create_usage_tracking_state(
|
||||
num_choices=1,
|
||||
reasoning_parser=None,
|
||||
enable_prompt_tokens_details=True,
|
||||
)
|
||||
state.num_cached_tokens = 0
|
||||
|
||||
chunk = {
|
||||
"id": "chatcmpl-test",
|
||||
"object": "chat.completion.chunk",
|
||||
"choices": [],
|
||||
"usage": {
|
||||
"prompt_tokens": 3,
|
||||
"completion_tokens": 4,
|
||||
"total_tokens": 7,
|
||||
},
|
||||
}
|
||||
|
||||
data = usage_patch._inject_stream_usage_details(
|
||||
f"data: {json.dumps(chunk)}\n\n",
|
||||
state,
|
||||
)
|
||||
payload = json.loads(data.removeprefix("data: ").removesuffix("\n\n"))
|
||||
|
||||
assert payload["usage"]["prompt_tokens_details"] == {
|
||||
"cached_tokens": 0,
|
||||
}
|
||||
assert "completion_tokens_details" not in payload["usage"]
|
||||
57
tests/ut/patch/platform/test_patch_pp_mtp.py
Normal file
57
tests/ut/patch/platform/test_patch_pp_mtp.py
Normal file
@@ -0,0 +1,57 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
|
||||
from types import SimpleNamespace
|
||||
|
||||
import pytest
|
||||
from vllm.config.model import ModelConfig
|
||||
|
||||
|
||||
def test_model_config_validates_local_mtp_drafter_as_single_pp_rank(monkeypatch):
|
||||
fake_registry = SimpleNamespace(
|
||||
is_pp_supported_model=lambda _architectures, _model_config: False,
|
||||
)
|
||||
monkeypatch.setattr(ModelConfig, "registry", property(lambda _self: fake_registry))
|
||||
|
||||
model_config = ModelConfig.__new__(ModelConfig)
|
||||
model_config.hf_config = SimpleNamespace(model_type="qwen3_5_mtp")
|
||||
model_config.runner = "draft"
|
||||
model_config.model_arch_config = SimpleNamespace(
|
||||
total_num_attention_heads=1,
|
||||
architectures=["Qwen3_5MTP"],
|
||||
)
|
||||
model_config.multimodal_config = None
|
||||
|
||||
parallel_config = SimpleNamespace(
|
||||
tensor_parallel_size=1,
|
||||
enable_expert_parallel=False,
|
||||
pipeline_parallel_size=2,
|
||||
decode_context_parallel_size=1,
|
||||
)
|
||||
|
||||
ModelConfig.verify_with_parallel_config(model_config, parallel_config)
|
||||
assert parallel_config.pipeline_parallel_size == 2
|
||||
|
||||
|
||||
def test_model_config_keeps_target_model_pp_validation(monkeypatch):
|
||||
fake_registry = SimpleNamespace(
|
||||
is_pp_supported_model=lambda _architectures, _model_config: False,
|
||||
)
|
||||
monkeypatch.setattr(ModelConfig, "registry", property(lambda _self: fake_registry))
|
||||
|
||||
model_config = ModelConfig.__new__(ModelConfig)
|
||||
model_config.hf_config = SimpleNamespace(model_type="qwen3_5_mtp")
|
||||
model_config.runner = "generate"
|
||||
model_config.model_arch_config = SimpleNamespace(
|
||||
total_num_attention_heads=1,
|
||||
architectures=["UnsupportedForPP"],
|
||||
)
|
||||
|
||||
parallel_config = SimpleNamespace(
|
||||
tensor_parallel_size=1,
|
||||
enable_expert_parallel=False,
|
||||
pipeline_parallel_size=2,
|
||||
decode_context_parallel_size=1,
|
||||
)
|
||||
|
||||
with pytest.raises(NotImplementedError):
|
||||
ModelConfig.verify_with_parallel_config(model_config, parallel_config)
|
||||
120
tests/ut/patch/platform/test_patch_shm_broadcast.py
Normal file
120
tests/ut/patch/platform/test_patch_shm_broadcast.py
Normal file
@@ -0,0 +1,120 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
|
||||
import threading
|
||||
from unittest import mock
|
||||
|
||||
import pytest
|
||||
from vllm.distributed.device_communicators.shm_broadcast import MessageQueue
|
||||
|
||||
import vllm_ascend.patch.platform.patch_shm_broadcast as patch
|
||||
|
||||
|
||||
@pytest.mark.parametrize("should_warn", [False, True])
|
||||
def test_reader_timeout_caps_indefinite_waits(monkeypatch, should_warn):
|
||||
monkeypatch.setattr(patch, "SHM_READER_RECHECK_INTERVAL_MS", 7)
|
||||
timeout = MessageQueue.ReadTimeoutWithWarnings(timeout=None, should_warn=should_warn)
|
||||
assert timeout.timeout_ms() == 7
|
||||
|
||||
|
||||
def test_reader_rechecks_shm_after_lost_notify(monkeypatch):
|
||||
monkeypatch.setattr(patch, "SHM_READER_RECHECK_INTERVAL_MS", 50)
|
||||
writer = MessageQueue(
|
||||
n_reader=1,
|
||||
n_local_reader=1,
|
||||
max_chunk_bytes=1024 * 1024,
|
||||
max_chunks=1,
|
||||
)
|
||||
reader = MessageQueue.create_from_handle(writer.export_handle(), rank=0)
|
||||
poll_started = threading.Event()
|
||||
allow_timeout = threading.Event()
|
||||
result = {}
|
||||
|
||||
def acquire_read():
|
||||
try:
|
||||
with reader.acquire_read(indefinite=True) as buf:
|
||||
result["value"] = buf[0]
|
||||
except Exception as exc:
|
||||
result["exception"] = exc
|
||||
|
||||
def poll_timeout(*, timeout: int | None = None):
|
||||
poll_started.set()
|
||||
assert allow_timeout.wait(timeout=5)
|
||||
return []
|
||||
|
||||
try:
|
||||
writer.wait_until_ready()
|
||||
reader.wait_until_ready()
|
||||
reader._spin_condition.last_read = 0
|
||||
reader._spin_condition.busy_loop_s = 0
|
||||
|
||||
with mock.patch.object(
|
||||
reader._spin_condition.poller,
|
||||
"poll",
|
||||
side_effect=poll_timeout,
|
||||
) as poll:
|
||||
read_thread = threading.Thread(target=acquire_read, daemon=True)
|
||||
read_thread.start()
|
||||
assert poll_started.wait(timeout=5)
|
||||
with writer.acquire_write(timeout=0.1) as buf:
|
||||
buf[0] = 123
|
||||
allow_timeout.set()
|
||||
read_thread.join(timeout=5)
|
||||
|
||||
assert not read_thread.is_alive()
|
||||
poll.assert_called_once_with(timeout=50)
|
||||
|
||||
if exception := result.get("exception"):
|
||||
raise exception
|
||||
assert result["value"] == 123
|
||||
finally:
|
||||
writer.shutdown()
|
||||
reader.shutdown()
|
||||
for socket in (
|
||||
writer.local_socket,
|
||||
writer._spin_condition.local_notify_socket,
|
||||
reader.local_socket,
|
||||
reader._spin_condition.local_notify_socket,
|
||||
reader._spin_condition.read_cancel_socket,
|
||||
reader._spin_condition.write_cancel_socket,
|
||||
):
|
||||
socket.close(linger=0)
|
||||
|
||||
|
||||
def test_acquire_read_releases_slot_when_reader_raises():
|
||||
writer = MessageQueue(
|
||||
n_reader=1,
|
||||
n_local_reader=1,
|
||||
max_chunk_bytes=1024 * 1024,
|
||||
max_chunks=1,
|
||||
)
|
||||
reader = MessageQueue.create_from_handle(writer.export_handle(), rank=0)
|
||||
try:
|
||||
writer.wait_until_ready()
|
||||
reader.wait_until_ready()
|
||||
writer.enqueue({"payload": "first"})
|
||||
|
||||
with (
|
||||
pytest.raises(RuntimeError, match="reader failed"),
|
||||
reader.acquire_read(timeout=0.1),
|
||||
):
|
||||
raise RuntimeError("reader failed")
|
||||
|
||||
with writer.buffer.get_metadata(0) as metadata_buffer:
|
||||
assert metadata_buffer[0] == 1
|
||||
assert metadata_buffer[1] == 1
|
||||
|
||||
with writer.acquire_write(timeout=0.1) as buf:
|
||||
buf[0] = 0
|
||||
finally:
|
||||
writer.shutdown()
|
||||
reader.shutdown()
|
||||
for socket in (
|
||||
writer.local_socket,
|
||||
writer._spin_condition.local_notify_socket,
|
||||
reader.local_socket,
|
||||
reader._spin_condition.local_notify_socket,
|
||||
reader._spin_condition.read_cancel_socket,
|
||||
reader._spin_condition.write_cancel_socket,
|
||||
):
|
||||
socket.close(linger=0)
|
||||
215
tests/ut/patch/platform/test_patch_structured_output.py
Normal file
215
tests/ut/patch/platform/test_patch_structured_output.py
Normal file
@@ -0,0 +1,215 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
|
||||
from inspect import signature
|
||||
from types import SimpleNamespace
|
||||
|
||||
import pytest
|
||||
import vllm.v1.structured_output as structured_output
|
||||
from vllm.config.structured_outputs import StructuredOutputsConfig
|
||||
from vllm.sampling_params import SamplingParams, StructuredOutputsParams
|
||||
from vllm.v1.structured_output import StructuredOutputManager, backend_guidance, backend_xgrammar
|
||||
from vllm.v1.structured_output.backend_types import StructuredOutputOptions
|
||||
|
||||
from vllm_ascend.patch.platform import patch_structured_output # noqa: F401
|
||||
|
||||
MODEL_CONFIG = SimpleNamespace(is_diffusion=False)
|
||||
|
||||
|
||||
class FakeBackend:
|
||||
def __init__(self, vllm_config, tokenizer, vocab_size):
|
||||
self.vllm_config = vllm_config
|
||||
self.tokenizer = tokenizer
|
||||
self.vocab_size = vocab_size
|
||||
|
||||
def compile_grammar(self, request_type, grammar_spec):
|
||||
return (type(self).__name__, request_type, grammar_spec)
|
||||
|
||||
|
||||
class FakeXgrammarBackend(FakeBackend):
|
||||
pass
|
||||
|
||||
|
||||
class FakeGuidanceBackend(FakeBackend):
|
||||
pass
|
||||
|
||||
|
||||
def make_manager() -> StructuredOutputManager:
|
||||
manager = object.__new__(StructuredOutputManager)
|
||||
manager.backend = None
|
||||
manager.vllm_config = SimpleNamespace(model_config=SimpleNamespace(get_vocab_size=lambda: 128))
|
||||
manager.tokenizer = object()
|
||||
manager._use_async_grammar_compilation = False
|
||||
return manager
|
||||
|
||||
|
||||
def make_request(backend: str):
|
||||
return SimpleNamespace(
|
||||
sampling_params=SimpleNamespace(structured_outputs=SimpleNamespace(_backend=backend)),
|
||||
structured_output_request=SimpleNamespace(
|
||||
structured_output_key=(StructuredOutputOptions.JSON, "{}"),
|
||||
grammar=None,
|
||||
),
|
||||
)
|
||||
|
||||
|
||||
def validate_structured_outputs(params, config):
|
||||
original_validate = getattr(
|
||||
SamplingParams,
|
||||
patch_structured_output._ORIGINAL_VALIDATE_ATTR,
|
||||
)
|
||||
if "model_config" in signature(original_validate).parameters:
|
||||
params._validate_structured_outputs(MODEL_CONFIG, config, tokenizer=object())
|
||||
else:
|
||||
params._validate_structured_outputs(config, tokenizer=object())
|
||||
|
||||
|
||||
def test_sampling_params_rejects_mixed_structured_output_backends(monkeypatch):
|
||||
def fake_validate_xgrammar(sampling_params):
|
||||
schema = sampling_params.structured_outputs.json
|
||||
if schema.get("force_guidance"):
|
||||
raise ValueError("xgrammar unsupported")
|
||||
|
||||
monkeypatch.setattr(
|
||||
backend_xgrammar,
|
||||
"validate_xgrammar_grammar",
|
||||
fake_validate_xgrammar,
|
||||
)
|
||||
monkeypatch.setattr(
|
||||
backend_guidance,
|
||||
"has_guidance_unsupported_json_features",
|
||||
lambda schema: False,
|
||||
)
|
||||
monkeypatch.setattr(
|
||||
backend_guidance,
|
||||
"validate_guidance_grammar",
|
||||
lambda sampling_params, tokenizer=None: None,
|
||||
)
|
||||
|
||||
config = StructuredOutputsConfig(backend="auto")
|
||||
xgrammar_params = SamplingParams(structured_outputs=StructuredOutputsParams(json={"type": "object"}))
|
||||
validate_structured_outputs(xgrammar_params, config)
|
||||
|
||||
assert xgrammar_params.structured_outputs._backend == "xgrammar"
|
||||
assert getattr(config, patch_structured_output._BACKEND_ATTR) == "xgrammar"
|
||||
|
||||
guidance_params = SamplingParams(structured_outputs=StructuredOutputsParams(json={"force_guidance": True}))
|
||||
with pytest.raises(ValueError, match="already using 'xgrammar'.*'guidance'"):
|
||||
validate_structured_outputs(guidance_params, config)
|
||||
|
||||
|
||||
def test_sampling_params_allows_consistent_guidance_backend(monkeypatch):
|
||||
monkeypatch.setattr(
|
||||
backend_guidance,
|
||||
"has_guidance_unsupported_json_features",
|
||||
lambda schema: False,
|
||||
)
|
||||
monkeypatch.setattr(
|
||||
backend_guidance,
|
||||
"validate_guidance_grammar",
|
||||
lambda sampling_params, tokenizer=None: None,
|
||||
)
|
||||
|
||||
config = StructuredOutputsConfig(backend="guidance")
|
||||
for _ in range(2):
|
||||
params = SamplingParams(structured_outputs=StructuredOutputsParams(json={"type": "array"}))
|
||||
validate_structured_outputs(params, config)
|
||||
|
||||
assert params.structured_outputs._backend == "guidance"
|
||||
assert getattr(config, patch_structured_output._BACKEND_ATTR) == "guidance"
|
||||
|
||||
|
||||
def test_failed_first_validation_does_not_lock_config(monkeypatch):
|
||||
monkeypatch.setattr(
|
||||
backend_xgrammar,
|
||||
"validate_xgrammar_grammar",
|
||||
lambda sampling_params: (_ for _ in ()).throw(ValueError("xgrammar error")),
|
||||
)
|
||||
monkeypatch.setattr(
|
||||
backend_guidance,
|
||||
"has_guidance_unsupported_json_features",
|
||||
lambda schema: False,
|
||||
)
|
||||
monkeypatch.setattr(
|
||||
backend_guidance,
|
||||
"validate_guidance_grammar",
|
||||
lambda sampling_params, tokenizer=None: (_ for _ in ()).throw(ValueError("guidance error")),
|
||||
)
|
||||
|
||||
config = StructuredOutputsConfig(backend="auto")
|
||||
params = SamplingParams(structured_outputs=StructuredOutputsParams(json={"force_guidance": True}))
|
||||
with pytest.raises(ValueError, match="guidance error"):
|
||||
validate_structured_outputs(params, config)
|
||||
|
||||
assert not hasattr(config, patch_structured_output._BACKEND_ATTR)
|
||||
|
||||
|
||||
def test_manager_rejects_mixed_structured_output_backends(monkeypatch):
|
||||
monkeypatch.setattr(structured_output, "XgrammarBackend", FakeXgrammarBackend)
|
||||
monkeypatch.setattr(structured_output, "GuidanceBackend", FakeGuidanceBackend)
|
||||
|
||||
manager = make_manager()
|
||||
xgrammar_request = make_request("xgrammar")
|
||||
manager.grammar_init(xgrammar_request)
|
||||
|
||||
assert isinstance(manager.backend, FakeXgrammarBackend)
|
||||
assert (
|
||||
getattr(
|
||||
manager,
|
||||
patch_structured_output._BACKEND_ATTR,
|
||||
)
|
||||
== "xgrammar"
|
||||
)
|
||||
assert xgrammar_request.structured_output_request.grammar == (
|
||||
"FakeXgrammarBackend",
|
||||
StructuredOutputOptions.JSON,
|
||||
"{}",
|
||||
)
|
||||
|
||||
guidance_request = make_request("guidance")
|
||||
with pytest.raises(ValueError, match="already using 'xgrammar'.*'guidance'"):
|
||||
manager.grammar_init(guidance_request)
|
||||
|
||||
|
||||
def test_manager_rejects_mixed_backend_after_subclassed_backend_is_initialized():
|
||||
manager = make_manager()
|
||||
manager.backend = FakeXgrammarBackend(
|
||||
manager.vllm_config,
|
||||
manager.tokenizer,
|
||||
manager.vllm_config.model_config.get_vocab_size(),
|
||||
)
|
||||
|
||||
with pytest.raises(ValueError, match="already using 'xgrammar'.*'guidance'"):
|
||||
manager.grammar_init(make_request("guidance"))
|
||||
|
||||
|
||||
def test_manager_allows_consistent_guidance_backend(monkeypatch):
|
||||
monkeypatch.setattr(structured_output, "GuidanceBackend", FakeGuidanceBackend)
|
||||
|
||||
manager = make_manager()
|
||||
for _ in range(2):
|
||||
request = make_request("guidance")
|
||||
manager.grammar_init(request)
|
||||
|
||||
assert isinstance(manager.backend, FakeGuidanceBackend)
|
||||
assert getattr(manager, patch_structured_output._BACKEND_ATTR) == "guidance"
|
||||
assert request.structured_output_request.grammar == (
|
||||
"FakeGuidanceBackend",
|
||||
StructuredOutputOptions.JSON,
|
||||
"{}",
|
||||
)
|
||||
|
||||
|
||||
def test_failed_first_backend_does_not_lock_manager(monkeypatch):
|
||||
monkeypatch.setattr(structured_output, "XgrammarBackend", FakeXgrammarBackend)
|
||||
|
||||
manager = make_manager()
|
||||
with pytest.raises(ValueError, match="Unsupported structured output backend"):
|
||||
manager.grammar_init(make_request("unsupported"))
|
||||
|
||||
assert not hasattr(manager, patch_structured_output._BACKEND_ATTR)
|
||||
|
||||
request = make_request("xgrammar")
|
||||
manager.grammar_init(request)
|
||||
|
||||
assert isinstance(manager.backend, FakeXgrammarBackend)
|
||||
assert getattr(manager, patch_structured_output._BACKEND_ATTR) == "xgrammar"
|
||||
194
tests/ut/patch/platform/test_patch_tool_choice_none_content.py
Normal file
194
tests/ut/patch/platform/test_patch_tool_choice_none_content.py
Normal file
@@ -0,0 +1,194 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
|
||||
from openai.types.chat.chat_completion import ChatCompletion as OpenAIChatCompletion
|
||||
from openai.types.chat.chat_completion_chunk import ChatCompletionChunk
|
||||
from vllm.entrypoints.openai.chat_completion.protocol import (
|
||||
ChatCompletionResponse,
|
||||
ChatCompletionResponseChoice,
|
||||
ChatCompletionResponseStreamChoice,
|
||||
ChatCompletionStreamResponse,
|
||||
ChatMessage,
|
||||
)
|
||||
from vllm.entrypoints.openai.engine.protocol import (
|
||||
DeltaFunctionCall,
|
||||
DeltaMessage,
|
||||
DeltaToolCall,
|
||||
FunctionCall,
|
||||
ToolCall,
|
||||
UsageInfo,
|
||||
)
|
||||
from vllm.entrypoints.openai.responses.protocol import ResponsesRequest
|
||||
from vllm.parser.abstract_parser import DelegatingParser
|
||||
|
||||
from vllm_ascend.patch.platform import patch_tool_choice_none_content # noqa: F401
|
||||
|
||||
|
||||
class _DummyDelegatingParser(DelegatingParser):
|
||||
def is_reasoning_end(self, input_ids: list[int]) -> bool:
|
||||
return False
|
||||
|
||||
def extract_content_ids(self, input_ids: list[int]) -> list[int]:
|
||||
return input_ids
|
||||
|
||||
def extract_reasoning(self, model_output: str, request):
|
||||
return None, model_output
|
||||
|
||||
def extract_reasoning_streaming(
|
||||
self,
|
||||
previous_text: str,
|
||||
current_text: str,
|
||||
delta_text: str,
|
||||
previous_token_ids: list[int],
|
||||
current_token_ids: list[int],
|
||||
delta_token_ids: list[int],
|
||||
):
|
||||
return None
|
||||
|
||||
def extract_tool_calls(self, model_output: str, request):
|
||||
return None
|
||||
|
||||
|
||||
def test_responses_parser_allows_named_tool_choice_with_none_content():
|
||||
request = ResponsesRequest.model_validate(
|
||||
{
|
||||
"model": "test-model",
|
||||
"input": "test",
|
||||
"tools": [
|
||||
{
|
||||
"type": "function",
|
||||
"name": "get_weather",
|
||||
"parameters": {"type": "object", "properties": {}},
|
||||
}
|
||||
],
|
||||
"tool_choice": {"type": "function", "name": "get_weather"},
|
||||
}
|
||||
)
|
||||
parser = _DummyDelegatingParser(tokenizer=None)
|
||||
|
||||
tool_calls, content = parser._extract_tool_calls(
|
||||
content=None,
|
||||
request=request,
|
||||
enable_auto_tools=False,
|
||||
)
|
||||
|
||||
assert content is None
|
||||
assert tool_calls == []
|
||||
|
||||
|
||||
def _chat_response(message: ChatMessage) -> ChatCompletionResponse:
|
||||
return ChatCompletionResponse(
|
||||
model="test-model",
|
||||
choices=[
|
||||
ChatCompletionResponseChoice(
|
||||
index=0,
|
||||
message=message,
|
||||
finish_reason="stop",
|
||||
)
|
||||
],
|
||||
usage=UsageInfo(prompt_tokens=1, completion_tokens=1, total_tokens=2),
|
||||
)
|
||||
|
||||
|
||||
def test_chat_completion_response_omits_empty_tool_calls_payload():
|
||||
response = _chat_response(ChatMessage(role="assistant", content="done"))
|
||||
|
||||
payload = response.model_dump()
|
||||
payload_json = response.model_dump_json()
|
||||
|
||||
assert "tool_calls" not in payload["choices"][0]["message"]
|
||||
parsed = OpenAIChatCompletion.model_validate(payload)
|
||||
assert parsed.choices[0].message.tool_calls is None
|
||||
parsed_json = OpenAIChatCompletion.model_validate_json(payload_json)
|
||||
assert parsed_json.choices[0].message.tool_calls is None
|
||||
|
||||
|
||||
def test_chat_completion_response_model_dump_json_uses_json_mode(monkeypatch):
|
||||
seen_kwargs = {}
|
||||
|
||||
def fake_model_dump(self, *args, **kwargs):
|
||||
seen_kwargs.update(kwargs)
|
||||
return {"choices": [{"message": {"tool_calls": []}}]}
|
||||
|
||||
monkeypatch.setattr(
|
||||
patch_tool_choice_none_content,
|
||||
"_original_chat_completion_response_model_dump",
|
||||
fake_model_dump,
|
||||
)
|
||||
|
||||
response = _chat_response(ChatMessage(role="assistant", content="done"))
|
||||
payload_json = response.model_dump_json()
|
||||
|
||||
assert seen_kwargs["mode"] == "json"
|
||||
assert payload_json == '{"choices":[{"message":{}}]}'
|
||||
|
||||
|
||||
def test_chat_completion_response_keeps_non_empty_tool_calls_payload():
|
||||
response = _chat_response(
|
||||
ChatMessage(
|
||||
role="assistant",
|
||||
content="",
|
||||
tool_calls=[
|
||||
ToolCall(
|
||||
function=FunctionCall(
|
||||
name="get_weather",
|
||||
arguments='{"city": "Beijing"}',
|
||||
)
|
||||
)
|
||||
],
|
||||
)
|
||||
)
|
||||
|
||||
message = response.model_dump()["choices"][0]["message"]
|
||||
|
||||
assert len(message["tool_calls"]) == 1
|
||||
assert message["tool_calls"][0]["function"]["name"] == "get_weather"
|
||||
|
||||
|
||||
def _stream_response(delta: DeltaMessage) -> ChatCompletionStreamResponse:
|
||||
return ChatCompletionStreamResponse(
|
||||
id="chatcmpl-test",
|
||||
object="chat.completion.chunk",
|
||||
created=1,
|
||||
model="test-model",
|
||||
choices=[
|
||||
ChatCompletionResponseStreamChoice(
|
||||
index=0,
|
||||
delta=delta,
|
||||
finish_reason=None,
|
||||
)
|
||||
],
|
||||
)
|
||||
|
||||
|
||||
def test_chat_completion_stream_response_omits_empty_tool_calls_payload():
|
||||
response = _stream_response(DeltaMessage(content="done", tool_calls=[]))
|
||||
|
||||
payload = response.model_dump(exclude_unset=True)
|
||||
payload_json = response.model_dump_json(exclude_unset=True)
|
||||
|
||||
assert "tool_calls" not in payload["choices"][0]["delta"]
|
||||
parsed = ChatCompletionChunk.model_validate_json(payload_json)
|
||||
assert parsed.choices[0].delta.tool_calls is None
|
||||
|
||||
|
||||
def test_chat_completion_stream_response_keeps_non_empty_tool_calls_payload():
|
||||
response = _stream_response(
|
||||
DeltaMessage(
|
||||
tool_calls=[
|
||||
DeltaToolCall(
|
||||
index=0,
|
||||
id="call-test",
|
||||
type="function",
|
||||
function=DeltaFunctionCall(
|
||||
name="get_weather",
|
||||
arguments='{"city": "Beijing"}',
|
||||
),
|
||||
)
|
||||
]
|
||||
)
|
||||
)
|
||||
|
||||
delta = response.model_dump(exclude_unset=True)["choices"][0]["delta"]
|
||||
|
||||
assert len(delta["tool_calls"]) == 1
|
||||
assert delta["tool_calls"][0]["function"]["name"] == "get_weather"
|
||||
523
tests/ut/patch/platform/test_prefix_cache_cp_patches.py
Normal file
523
tests/ut/patch/platform/test_prefix_cache_cp_patches.py
Normal file
@@ -0,0 +1,523 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
|
||||
import math
|
||||
from types import SimpleNamespace
|
||||
from unittest.mock import MagicMock
|
||||
|
||||
import pytest
|
||||
import torch
|
||||
from vllm.v1.core.block_pool import BlockPool
|
||||
from vllm.v1.core.single_type_kv_cache_manager import (
|
||||
SlidingWindowManager,
|
||||
)
|
||||
from vllm.v1.kv_cache_interface import (
|
||||
FullAttentionSpec,
|
||||
KVCacheConfig,
|
||||
KVCacheGroupSpec,
|
||||
KVCacheTensor,
|
||||
MambaSpec,
|
||||
MLAAttentionSpec,
|
||||
SlidingWindowMLASpec,
|
||||
UniformTypeKVCacheSpecs,
|
||||
)
|
||||
|
||||
from vllm_ascend.patch.platform.patch_kv_cache_coordinator import (
|
||||
AscendHybridKVCacheCoordinator,
|
||||
_is_deepseek_v4_kv_cache_spec,
|
||||
get_kv_cache_coordinator,
|
||||
)
|
||||
from vllm_ascend.patch.platform.patch_kv_cache_utils import (
|
||||
_ascend_resolve_kv_cache_block_sizes,
|
||||
)
|
||||
from vllm_ascend.patch.platform.patch_mamba_manager import AscendMambaManager
|
||||
|
||||
|
||||
def _make_hybrid_kv_cache_config(
|
||||
full_block_size: int = 16,
|
||||
mamba_block_size: int = 16,
|
||||
) -> KVCacheConfig:
|
||||
full_spec = FullAttentionSpec(
|
||||
block_size=full_block_size,
|
||||
num_kv_heads=8,
|
||||
head_size=64,
|
||||
dtype=torch.float16,
|
||||
)
|
||||
mamba_spec = MambaSpec(
|
||||
block_size=mamba_block_size,
|
||||
shapes=((1,),),
|
||||
dtypes=(torch.float32,),
|
||||
mamba_cache_mode="none",
|
||||
)
|
||||
return KVCacheConfig(
|
||||
num_blocks=10,
|
||||
kv_cache_tensors=[
|
||||
KVCacheTensor(size=full_spec.page_size_bytes * 10, shared_by=["attn"]),
|
||||
KVCacheTensor(size=mamba_spec.page_size_bytes * 10, shared_by=["mamba"]),
|
||||
],
|
||||
kv_cache_groups=[
|
||||
KVCacheGroupSpec(layer_names=["attn"], kv_cache_spec=full_spec),
|
||||
KVCacheGroupSpec(layer_names=["mamba"], kv_cache_spec=mamba_spec),
|
||||
],
|
||||
)
|
||||
|
||||
|
||||
def _make_deepseek_v4_kv_cache_config() -> KVCacheConfig:
|
||||
c4_spec = MLAAttentionSpec(
|
||||
block_size=128,
|
||||
num_kv_heads=1,
|
||||
head_size=128,
|
||||
dtype=torch.float16,
|
||||
compress_ratio=4,
|
||||
model_version="deepseek_v4",
|
||||
)
|
||||
c128_spec = MLAAttentionSpec(
|
||||
block_size=128,
|
||||
num_kv_heads=1,
|
||||
head_size=128,
|
||||
dtype=torch.float16,
|
||||
compress_ratio=128,
|
||||
model_version="deepseek_v4",
|
||||
)
|
||||
c4_group_spec = UniformTypeKVCacheSpecs.from_specs({"c4_attn": c4_spec})
|
||||
c128_group_spec = UniformTypeKVCacheSpecs.from_specs({"c128_attn": c128_spec})
|
||||
assert c4_group_spec is not None
|
||||
assert c128_group_spec is not None
|
||||
return KVCacheConfig(
|
||||
num_blocks=10,
|
||||
kv_cache_tensors=[
|
||||
KVCacheTensor(size=c4_spec.page_size_bytes * 10, shared_by=["c4_attn"]),
|
||||
KVCacheTensor(size=c128_spec.page_size_bytes * 10, shared_by=["c128_attn"]),
|
||||
],
|
||||
kv_cache_groups=[
|
||||
KVCacheGroupSpec(layer_names=["c4_attn"], kv_cache_spec=c4_group_spec),
|
||||
KVCacheGroupSpec(layer_names=["c128_attn"], kv_cache_spec=c128_group_spec),
|
||||
],
|
||||
)
|
||||
|
||||
|
||||
def _make_vllm_config(
|
||||
*,
|
||||
enable_prefix_caching: bool,
|
||||
dcp: int,
|
||||
pcp: int,
|
||||
block_size: int = 16,
|
||||
) -> SimpleNamespace:
|
||||
return SimpleNamespace(
|
||||
cache_config=SimpleNamespace(
|
||||
block_size=block_size,
|
||||
enable_prefix_caching=enable_prefix_caching,
|
||||
),
|
||||
parallel_config=SimpleNamespace(
|
||||
decode_context_parallel_size=dcp,
|
||||
prefill_context_parallel_size=pcp,
|
||||
),
|
||||
)
|
||||
|
||||
|
||||
def _make_coordinator_for_effective_block_size(
|
||||
*,
|
||||
dcp_world_size: int,
|
||||
pcp_world_size: int,
|
||||
enable_caching: bool,
|
||||
) -> AscendHybridKVCacheCoordinator:
|
||||
coordinator = AscendHybridKVCacheCoordinator.__new__(AscendHybridKVCacheCoordinator)
|
||||
coordinator.dcp_world_size = dcp_world_size
|
||||
coordinator.pcp_world_size = pcp_world_size
|
||||
coordinator.enable_caching = enable_caching
|
||||
return coordinator
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("enable_prefix_caching", "expected_hash_block_size"),
|
||||
[
|
||||
pytest.param(False, math.lcm(16, 32) * 2 * 2, id="cp-without-prefix-caching"),
|
||||
pytest.param(True, math.gcd(16, 32), id="cp-with-prefix-caching"),
|
||||
],
|
||||
)
|
||||
def test_resolve_kv_cache_block_sizes_with_cp_hybrid_groups(
|
||||
enable_prefix_caching: bool,
|
||||
expected_hash_block_size: int,
|
||||
) -> None:
|
||||
kv_cache_config = _make_hybrid_kv_cache_config(full_block_size=16, mamba_block_size=32)
|
||||
vllm_config = _make_vllm_config(
|
||||
enable_prefix_caching=enable_prefix_caching,
|
||||
dcp=2,
|
||||
pcp=2,
|
||||
)
|
||||
|
||||
scheduler_block_size, hash_block_size = _ascend_resolve_kv_cache_block_sizes(
|
||||
kv_cache_config,
|
||||
vllm_config,
|
||||
)
|
||||
|
||||
expected_scheduler_block_size = math.lcm(16, 32) * 2 * 2
|
||||
assert scheduler_block_size == expected_scheduler_block_size
|
||||
assert hash_block_size == expected_hash_block_size
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("spec_factory", "dcp", "pcp", "enable_caching", "expected"),
|
||||
[
|
||||
pytest.param(
|
||||
lambda: FullAttentionSpec(
|
||||
block_size=16,
|
||||
num_kv_heads=8,
|
||||
head_size=64,
|
||||
dtype=torch.float16,
|
||||
),
|
||||
2,
|
||||
2,
|
||||
True,
|
||||
64,
|
||||
id="full-attention-scales-with-cp",
|
||||
),
|
||||
pytest.param(
|
||||
lambda: MambaSpec(
|
||||
block_size=16,
|
||||
shapes=((1,),),
|
||||
dtypes=(torch.float32,),
|
||||
mamba_cache_mode="none",
|
||||
),
|
||||
2,
|
||||
2,
|
||||
True,
|
||||
16,
|
||||
id="mamba-keeps-physical-block-size-with-prefix-caching",
|
||||
),
|
||||
pytest.param(
|
||||
lambda: FullAttentionSpec(
|
||||
block_size=16,
|
||||
num_kv_heads=8,
|
||||
head_size=64,
|
||||
dtype=torch.float16,
|
||||
),
|
||||
1,
|
||||
1,
|
||||
True,
|
||||
16,
|
||||
id="full-attention-no-cp",
|
||||
),
|
||||
],
|
||||
)
|
||||
def test_get_effective_block_size(
|
||||
spec_factory,
|
||||
dcp: int,
|
||||
pcp: int,
|
||||
enable_caching: bool,
|
||||
expected: int,
|
||||
) -> None:
|
||||
coordinator = _make_coordinator_for_effective_block_size(
|
||||
dcp_world_size=dcp,
|
||||
pcp_world_size=pcp,
|
||||
enable_caching=enable_caching,
|
||||
)
|
||||
|
||||
assert coordinator._get_effective_block_size(spec_factory()) == expected
|
||||
|
||||
|
||||
def test_get_kv_cache_coordinator_delegates_single_group(monkeypatch) -> None:
|
||||
sentinel = object()
|
||||
kv_cache_config = _make_hybrid_kv_cache_config(full_block_size=16, mamba_block_size=16)
|
||||
single_group_config = KVCacheConfig(
|
||||
num_blocks=kv_cache_config.num_blocks,
|
||||
kv_cache_tensors=kv_cache_config.kv_cache_tensors[:1],
|
||||
kv_cache_groups=kv_cache_config.kv_cache_groups[:1],
|
||||
)
|
||||
|
||||
def _fake_orig(*args, **kwargs):
|
||||
return sentinel
|
||||
|
||||
monkeypatch.setattr(
|
||||
"vllm_ascend.patch.platform.patch_kv_cache_coordinator._orig_get_kv_cache_coordinator",
|
||||
_fake_orig,
|
||||
)
|
||||
|
||||
coordinator = get_kv_cache_coordinator(
|
||||
single_group_config,
|
||||
max_model_len=1024,
|
||||
max_num_batched_tokens=1024,
|
||||
use_eagle=False,
|
||||
enable_caching=True,
|
||||
enable_kv_cache_events=False,
|
||||
dcp_world_size=1,
|
||||
pcp_world_size=1,
|
||||
hash_block_size=16,
|
||||
)
|
||||
|
||||
assert coordinator is sentinel
|
||||
|
||||
|
||||
def test_get_kv_cache_coordinator_delegates_hybrid_without_caching(monkeypatch) -> None:
|
||||
sentinel = object()
|
||||
kv_cache_config = _make_hybrid_kv_cache_config(full_block_size=16, mamba_block_size=16)
|
||||
|
||||
def _fake_orig(*args, **kwargs):
|
||||
return sentinel
|
||||
|
||||
monkeypatch.setattr(
|
||||
"vllm_ascend.patch.platform.patch_kv_cache_coordinator._orig_get_kv_cache_coordinator",
|
||||
_fake_orig,
|
||||
)
|
||||
|
||||
coordinator = get_kv_cache_coordinator(
|
||||
kv_cache_config,
|
||||
max_model_len=1024,
|
||||
max_num_batched_tokens=1024,
|
||||
use_eagle=False,
|
||||
enable_caching=False,
|
||||
enable_kv_cache_events=False,
|
||||
dcp_world_size=2,
|
||||
pcp_world_size=2,
|
||||
hash_block_size=16,
|
||||
)
|
||||
|
||||
assert coordinator is sentinel
|
||||
|
||||
|
||||
def test_get_kv_cache_coordinator_uses_ascend_for_deepseek_v4(monkeypatch) -> None:
|
||||
sentinel = object()
|
||||
kv_cache_config = _make_deepseek_v4_kv_cache_config()
|
||||
|
||||
def _fake_orig(*args, **kwargs):
|
||||
raise AssertionError("DeepSeek V4 should use AscendHybridKVCacheCoordinator")
|
||||
|
||||
def _fake_ascend_coordinator(*args, **kwargs):
|
||||
return sentinel
|
||||
|
||||
monkeypatch.setattr(
|
||||
"vllm_ascend.patch.platform.patch_kv_cache_coordinator._orig_get_kv_cache_coordinator",
|
||||
_fake_orig,
|
||||
)
|
||||
monkeypatch.setattr(
|
||||
"vllm_ascend.patch.platform.patch_kv_cache_coordinator.AscendHybridKVCacheCoordinator",
|
||||
_fake_ascend_coordinator,
|
||||
)
|
||||
|
||||
coordinator = get_kv_cache_coordinator(
|
||||
kv_cache_config,
|
||||
max_model_len=1024,
|
||||
max_num_batched_tokens=1024,
|
||||
use_eagle=False,
|
||||
enable_caching=True,
|
||||
enable_kv_cache_events=False,
|
||||
dcp_world_size=1,
|
||||
pcp_world_size=1,
|
||||
hash_block_size=128,
|
||||
)
|
||||
|
||||
assert coordinator is sentinel
|
||||
|
||||
|
||||
class _FakeEagleManager:
|
||||
def __init__(self) -> None:
|
||||
self.use_eagle = False
|
||||
|
||||
|
||||
def test_verify_and_split_propagates_eagle_to_managers() -> None:
|
||||
"""Regression for DeepSeek-V4 prefix-cache hit rate 0% with MTP/EAGLE.
|
||||
|
||||
The eagle bit must reach each single-type manager: the SWA write path
|
||||
(``cache_blocks`` -> ``reachable_block_mask``) keys the retained checkpoint
|
||||
tail on ``manager.use_eagle``, while the read path
|
||||
(``find_longest_cache_hit``) applies ``drop_eagle_block`` to the same
|
||||
groups. If the manager keeps the default ``use_eagle=False`` the retained
|
||||
tail is one block short of the eagle peek boundary, the SWA group never
|
||||
hits, and the min-over-groups hybrid hit collapses to 0%.
|
||||
"""
|
||||
kv_cache_config = _make_deepseek_v4_kv_cache_config()
|
||||
|
||||
coordinator = AscendHybridKVCacheCoordinator.__new__(AscendHybridKVCacheCoordinator)
|
||||
coordinator.kv_cache_config = kv_cache_config
|
||||
coordinator.dcp_world_size = 1
|
||||
coordinator.pcp_world_size = 1
|
||||
coordinator.enable_caching = True
|
||||
# The c128 group (index 1) carries the EAGLE/MTP layers.
|
||||
coordinator.eagle_group_ids = {1}
|
||||
|
||||
coordinator.single_type_managers = (_FakeEagleManager(), _FakeEagleManager())
|
||||
|
||||
coordinator.verify_and_split_kv_cache_groups()
|
||||
|
||||
assert coordinator.single_type_managers[1].use_eagle is True
|
||||
assert coordinator.single_type_managers[0].use_eagle is False
|
||||
|
||||
|
||||
def test_verify_and_split_propagates_eagle_to_merged_spec_siblings() -> None:
|
||||
"""Upstream ``_annotate_eagle_groups_deepseek_v4`` flags only the single
|
||||
group holding the MTP layer, but the read path merges same-spec groups and
|
||||
applies ``drop_eagle_block`` to the whole merged group. So every sibling
|
||||
sharing that spec must also get ``use_eagle=True`` on the write path, else
|
||||
``get_cached_block`` (which needs the block cached for *all* group ids)
|
||||
misses and the hit collapses to 0%.
|
||||
"""
|
||||
base_config = _make_deepseek_v4_kv_cache_config()
|
||||
# Reuse the c128 spec object so the two c128 groups compare equal and merge
|
||||
# into one attention group in verify_and_split.
|
||||
c128_group_spec = base_config.kv_cache_groups[1].kv_cache_spec
|
||||
kv_cache_config = KVCacheConfig(
|
||||
num_blocks=base_config.num_blocks,
|
||||
kv_cache_tensors=base_config.kv_cache_tensors,
|
||||
kv_cache_groups=[
|
||||
base_config.kv_cache_groups[0], # c4 -> gid 0 (distinct spec)
|
||||
base_config.kv_cache_groups[1], # c128 -> gid 1
|
||||
KVCacheGroupSpec(layer_names=["c128_attn_mtp"], kv_cache_spec=c128_group_spec), # gid 2
|
||||
],
|
||||
)
|
||||
|
||||
coordinator = AscendHybridKVCacheCoordinator.__new__(AscendHybridKVCacheCoordinator)
|
||||
coordinator.kv_cache_config = kv_cache_config
|
||||
coordinator.dcp_world_size = 1
|
||||
coordinator.pcp_world_size = 1
|
||||
coordinator.enable_caching = True
|
||||
# Only the MTP sibling (gid 2) is flagged, exactly as upstream does.
|
||||
coordinator.eagle_group_ids = {2}
|
||||
|
||||
coordinator.single_type_managers = (
|
||||
_FakeEagleManager(),
|
||||
_FakeEagleManager(),
|
||||
_FakeEagleManager(),
|
||||
)
|
||||
|
||||
coordinator.verify_and_split_kv_cache_groups()
|
||||
|
||||
# Both gid 1 and gid 2 share the c128 spec and merge, so both must be eagle.
|
||||
assert coordinator.single_type_managers[1].use_eagle is True
|
||||
assert coordinator.single_type_managers[2].use_eagle is True
|
||||
assert coordinator.single_type_managers[0].use_eagle is False
|
||||
|
||||
|
||||
def test_mamba_eagle_lookup_does_not_expand_hybrid_hit() -> None:
|
||||
"""Mamba finders do not drop the EAGLE lookahead block.
|
||||
|
||||
The coordinator must therefore keep the Mamba lookup capped at the hit
|
||||
length already established by full attention.
|
||||
"""
|
||||
kv_cache_config = _make_hybrid_kv_cache_config()
|
||||
full_spec = kv_cache_config.kv_cache_groups[0].kv_cache_spec
|
||||
mamba_spec = kv_cache_config.kv_cache_groups[1].kv_cache_spec
|
||||
|
||||
class _FullHitManager:
|
||||
@classmethod
|
||||
def find_longest_cache_hit(cls, **kwargs):
|
||||
return ([object(), object()],)
|
||||
|
||||
class _MambaHitManager:
|
||||
lookup_max_lengths: list[int] = []
|
||||
|
||||
@classmethod
|
||||
def find_longest_cache_hit(cls, **kwargs):
|
||||
max_length = kwargs["max_length"]
|
||||
cls.lookup_max_lengths.append(max_length)
|
||||
block_size = kwargs["kv_cache_spec"].block_size
|
||||
return ([object()] * (max_length // block_size),)
|
||||
|
||||
coordinator = AscendHybridKVCacheCoordinator.__new__(AscendHybridKVCacheCoordinator)
|
||||
coordinator.kv_cache_config = kv_cache_config
|
||||
coordinator.attention_groups = [
|
||||
(full_spec, [0], _FullHitManager),
|
||||
(mamba_spec, [1], _MambaHitManager),
|
||||
]
|
||||
coordinator.eagle_attn_group_indices = {1}
|
||||
coordinator.dcp_world_size = 1
|
||||
coordinator.pcp_world_size = 1
|
||||
coordinator.enable_caching = True
|
||||
coordinator.hash_block_size = 16
|
||||
coordinator.lcm_block_size = 16
|
||||
coordinator.block_pool = MagicMock()
|
||||
|
||||
hit_blocks, hit_length = coordinator.find_longest_cache_hit(
|
||||
block_hashes=[MagicMock(), MagicMock(), MagicMock()],
|
||||
max_cache_hit_length=48,
|
||||
)
|
||||
|
||||
assert _MambaHitManager.lookup_max_lengths == [32]
|
||||
assert [len(blocks) for blocks in hit_blocks] == [2, 2]
|
||||
assert hit_length == 32
|
||||
|
||||
|
||||
def test_deepseek_v4_detection_handles_non_mapping_nested_specs() -> None:
|
||||
kv_cache_spec = SimpleNamespace(
|
||||
kv_cache_specs=[
|
||||
SimpleNamespace(model_version="deepseek_v4"),
|
||||
]
|
||||
)
|
||||
unknown_spec = SimpleNamespace(kv_cache_specs=object())
|
||||
|
||||
assert _is_deepseek_v4_kv_cache_spec(kv_cache_spec)
|
||||
assert not _is_deepseek_v4_kv_cache_spec(unknown_spec)
|
||||
|
||||
|
||||
def test_ascend_mamba_manager_uses_logical_block_size_with_prefix_caching() -> None:
|
||||
mamba_spec = MambaSpec(
|
||||
block_size=16,
|
||||
shapes=((1,),),
|
||||
dtypes=(torch.float32,),
|
||||
mamba_cache_mode="none",
|
||||
)
|
||||
block_pool = BlockPool(
|
||||
10,
|
||||
True,
|
||||
16,
|
||||
False,
|
||||
MagicMock(),
|
||||
)
|
||||
|
||||
manager_kwargs = dict(
|
||||
kv_cache_spec=mamba_spec,
|
||||
block_pool=block_pool,
|
||||
enable_caching=True,
|
||||
kv_cache_group_id=1,
|
||||
dcp_world_size=2,
|
||||
pcp_world_size=2,
|
||||
)
|
||||
manager_kwargs["scheduler_block_size"] = mamba_spec.block_size
|
||||
manager = AscendMambaManager(**manager_kwargs)
|
||||
|
||||
assert manager.block_size == mamba_spec.block_size
|
||||
|
||||
|
||||
def test_swa_reachable_block_mask_sparse_with_lcm_alignment() -> None:
|
||||
"""Regression: when ``scheduler_block_size`` is aligned to ``lcm_block_size``
|
||||
(instead of the raw-block-size LCM), ``SlidingWindowManager.reachable_block_mask``
|
||||
must produce a sparse mask rather than returning ``None``.
|
||||
|
||||
Before the fix, ``alignment_tokens`` was the LCM of raw block_sizes (e.g. 32),
|
||||
making ``need >= per_segment`` always true for Ascend's SWA configuration and
|
||||
the mask returned ``None`` (cache everything). After the fix the alignment is
|
||||
``lcm_block_size`` (e.g. 4096), which is large enough that only the tail
|
||||
blocks within each segment need caching.
|
||||
"""
|
||||
spec = SlidingWindowMLASpec(
|
||||
block_size=32, # Ascend SWA block_size (--block-size 32)
|
||||
num_kv_heads=1,
|
||||
head_size=512,
|
||||
dtype=torch.float32,
|
||||
sliding_window=128, # DeepSeek V4 window
|
||||
compress_ratio=1,
|
||||
)
|
||||
alignment_tokens = 4096 # lcm_block_size
|
||||
|
||||
mask = SlidingWindowManager.reachable_block_mask(
|
||||
start_block=0,
|
||||
end_block=256, # 256 × 32 = 8192 tokens (2 × alignment_tokens)
|
||||
alignment_tokens=alignment_tokens,
|
||||
kv_cache_spec=spec,
|
||||
use_eagle=False,
|
||||
retention_interval=None,
|
||||
num_prompt_tokens=None,
|
||||
)
|
||||
|
||||
# Must produce a sparse mask, not None.
|
||||
assert mask is not None, "should produce sparse mask with lcm alignment"
|
||||
|
||||
true_blocks = sum(mask)
|
||||
|
||||
# need = cdiv(window−1, block_size) = cdiv(127, 32) = 4
|
||||
# per_segment = alignment_tokens // block_size = 4096 // 32 = 128
|
||||
# Each 128-block segment caches the last 4 blocks (= 0 % sparse padding).
|
||||
total_blocks = len(mask)
|
||||
expected = 4 * (total_blocks // 128)
|
||||
assert true_blocks == expected, (
|
||||
f"expected {expected} cached blocks ({4}/{128} per segment), got {true_blocks}/{total_blocks}"
|
||||
)
|
||||
assert true_blocks > 0 and true_blocks < total_blocks, f"mask should be sparse, got {true_blocks}/{total_blocks}"
|
||||
Reference in New Issue
Block a user