init v0.23.0

Signed-off-by: Sun Ruoxi <sunruoxi@4paradigm.com>

tests/ut/patch/platform/test_deepseek_v4_thinking.py

@@ -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

tests/ut/patch/platform/test_patch_async_swa_kv_lifetime.py

@@ -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

tests/ut/patch/platform/test_patch_balance_schedule.py

@@ -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)

tests/ut/patch/platform/test_patch_deepseek_v4_tool_call_parser.py

@@ -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
+    )

tests/ut/patch/platform/test_patch_glm47_tool_call_parser.py

@@ -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)) == {}

tests/ut/patch/platform/test_patch_glm_tool_call_streaming.py

@@ -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]

tests/ut/patch/platform/test_patch_minimax_m2_tool_call_parser.py

@@ -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)

tests/ut/patch/platform/test_patch_minimax_usage_accounting.py

@@ -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"]

tests/ut/patch/platform/test_patch_pp_mtp.py

@@ -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)

tests/ut/patch/platform/test_patch_shm_broadcast.py

@@ -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)

tests/ut/patch/platform/test_patch_structured_output.py

@@ -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"

tests/ut/patch/platform/test_patch_tool_choice_none_content.py

@@ -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"

tests/ut/patch/platform/test_prefix_cache_cp_patches.py

@@ -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}"

tests/ut/patch/worker/patch_common/test_hccl_pg_registry.py

@@ -0,0 +1,365 @@
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# This file is a part of the vllm-ascend project.
+#
+from __future__ import annotations
+
+import sys
+from contextlib import contextmanager
+from dataclasses import dataclass
+from datetime import timedelta
+from importlib.util import module_from_spec, spec_from_file_location
+from pathlib import Path
+from typing import Any
+from unittest.mock import MagicMock
+
+_MODULE_PATH = Path(__file__).resolve().parents[5] / "vllm_ascend/patch/worker/_hccl_pg_registry.py"
+_MODULE_NAME = "vllm_ascend.patch.worker._hccl_pg_registry"
+_SPEC = spec_from_file_location(_MODULE_NAME, str(_MODULE_PATH))
+if _SPEC is None:
+    raise RuntimeError("Failed to load _hccl_pg_registry module spec")
+
+_MODULE: Any = module_from_spec(_SPEC)
+sys.modules[_MODULE_NAME] = _MODULE
+_SPEC.loader.exec_module(_MODULE)  # type: ignore[union-attr]
+
+RegistryEntry = _MODULE.RegistryEntry
+HcclPgRegistry = _MODULE.HcclPgRegistry
+make_hccl_pg_key = _MODULE.make_hccl_pg_key
+
+
+@contextmanager
+def _patch_destroy_process_group(destroy_fn):
+    previous = _MODULE._destroy_process_group
+    _MODULE._destroy_process_group = destroy_fn
+    try:
+        yield
+    finally:
+        _MODULE._destroy_process_group = previous
+
+
+@contextmanager
+def _set_non_group_member_sentinel(sentinel: object):
+    previous = (_MODULE._NON_GROUP_MEMBER, _MODULE._NON_GROUP_MEMBER_SET)
+    _MODULE._NON_GROUP_MEMBER = sentinel
+    _MODULE._NON_GROUP_MEMBER_SET = True
+    try:
+        yield
+    finally:
+        _MODULE._NON_GROUP_MEMBER, _MODULE._NON_GROUP_MEMBER_SET = previous
+
+
+@dataclass
+class FakeOptions:
+    hccl_config: dict[str, int] | None = None
+
+
+@dataclass
+class FakeOptionsWithUnknown:
+    hccl_config: dict[str, int] | None = None
+    non_default_option: int = 7
+
+
+@dataclass
+class RealisticFakeHcclOptions:
+    backend: str = "hccl"
+    global_ranks_in_group: list[int] | tuple[int, ...] = ()
+    group_id: str = ""
+    group_name: str = ""
+    hccl_config: dict[str, int] | None = None
+    is_high_priority_stream: bool = False
+    op_timeout: object = timedelta(seconds=10)
+
+
+def test_make_hccl_pg_key_respects_rank_order_and_reuse_domain():
+    opts = FakeOptions(hccl_config={"hccl_buffer_size": 200})
+    key_a = make_hccl_pg_key([0, 1], "hccl", opts, reuse_domain="shared")
+    key_b = make_hccl_pg_key([1, 0], "hccl", opts, reuse_domain="shared")
+    key_c = make_hccl_pg_key([0, 1], "hccl", opts, reuse_domain="eplb")
+
+    assert key_a != key_b
+    assert key_a != key_c
+
+
+def test_make_hccl_pg_key_mapping_hccl_config_affects_distinct_keys():
+    key_a = make_hccl_pg_key(
+        [0, 1],
+        "hccl",
+        {"hccl_config": {"hccl_buffer_size": 200}},
+        reuse_domain="shared",
+    )
+    key_b = make_hccl_pg_key(
+        [0, 1],
+        "hccl",
+        {"hccl_config": {"hccl_buffer_size": 400}},
+        reuse_domain="shared",
+    )
+
+    assert key_a != key_b
+
+
+def test_make_hccl_pg_key_accepts_realistic_options_object_defaults():
+    key_a = make_hccl_pg_key(
+        [0, 1],
+        "hccl",
+        RealisticFakeHcclOptions(hccl_config={"hccl_buffer_size": 200}),
+        reuse_domain="shared",
+    )
+    key_b = make_hccl_pg_key(
+        [0, 1],
+        "hccl",
+        RealisticFakeHcclOptions(hccl_config={"hccl_buffer_size": 400}),
+        reuse_domain="shared",
+    )
+
+    assert key_a is not None
+    assert key_b is not None
+    assert key_a != key_b
+
+
+def test_make_hccl_pg_key_accepts_matching_global_ranks_in_group():
+    key = make_hccl_pg_key(
+        [0, 1],
+        "hccl",
+        RealisticFakeHcclOptions(
+            global_ranks_in_group=[0, 1],
+            hccl_config={"hccl_buffer_size": 200},
+        ),
+        reuse_domain="shared",
+    )
+
+    assert key is not None
+
+
+def test_make_hccl_pg_key_fails_closed_on_mismatched_global_ranks_in_group():
+    key = make_hccl_pg_key(
+        [0, 1],
+        "hccl",
+        RealisticFakeHcclOptions(
+            global_ranks_in_group=[1, 2],
+            hccl_config={"hccl_buffer_size": 200},
+        ),
+        reuse_domain="shared",
+    )
+
+    assert key is None
+
+
+def test_make_hccl_pg_key_ignores_runtime_populated_group_identity_fields():
+    key_a = make_hccl_pg_key(
+        [0, 1],
+        "hccl",
+        RealisticFakeHcclOptions(
+            global_ranks_in_group=[0, 1],
+            group_id="hccl_pg_1",
+            group_name="tp_auto",
+            hccl_config={"hccl_buffer_size": 200},
+        ),
+        reuse_domain="shared",
+    )
+    key_b = make_hccl_pg_key(
+        [0, 1],
+        "hccl",
+        RealisticFakeHcclOptions(
+            global_ranks_in_group=[0, 1],
+            group_id="hccl_pg_2",
+            group_name="world_auto",
+            hccl_config={"hccl_buffer_size": 200},
+        ),
+        reuse_domain="shared",
+    )
+
+    assert key_a is not None
+    assert key_a == key_b
+
+
+def test_make_hccl_pg_key_fails_closed_for_unknown_mapping_fields():
+    assert (
+        make_hccl_pg_key(
+            [0, 1],
+            "hccl",
+            {"hccl_config": {"hccl_buffer_size": 200}, "non_default_field": 7},
+            reuse_domain="shared",
+        )
+        is None
+    )
+
+
+def test_registry_release_only_destroys_real_pg_at_zero_refcount():
+    destroy_fn = MagicMock()
+    with _patch_destroy_process_group(destroy_fn):
+        registry = HcclPgRegistry()
+        pg = object()
+        key = make_hccl_pg_key([0, 1], "hccl", FakeOptions(), reuse_domain="shared")
+        registry._entries[key] = RegistryEntry(handle=pg, refcount=1)
+
+        assert registry.release(key) == pg
+        assert key not in registry._entries
+        destroy_fn.assert_called_once_with(pg)
+
+
+def test_release_of_non_group_member_only_drops_registry_entry():
+    destroy_fn = MagicMock()
+    sentinel = _MODULE._load_non_group_member_sentinel()
+    with _patch_destroy_process_group(destroy_fn), _set_non_group_member_sentinel(sentinel):
+        registry = HcclPgRegistry()
+        key = make_hccl_pg_key([0, 1], "hccl", FakeOptions(), reuse_domain="shared")
+        registry._entries[key] = RegistryEntry(handle=sentinel, refcount=1)
+
+        assert registry.release(key) is None
+        assert key not in registry._entries
+        destroy_fn.assert_not_called()
+
+
+def test_acquire_reuses_cached_handle_and_refcount():
+    registry = HcclPgRegistry()
+    create_fn = MagicMock(side_effect=[MagicMock(name="first")])
+    destroy_fn = MagicMock()
+    key = make_hccl_pg_key(
+        [0, 1],
+        "hccl",
+        FakeOptions(hccl_config={"hccl_buffer_size": 200}),
+        reuse_domain="shared",
+    )
+
+    with _patch_destroy_process_group(destroy_fn):
+        first = registry.acquire(
+            ranks=[0, 1],
+            backend="hccl",
+            pg_options=FakeOptions(hccl_config={"hccl_buffer_size": 200}),
+            reuse_domain="shared",
+            create_fn=create_fn,
+        )
+        second = registry.acquire(
+            ranks=[0, 1],
+            backend="hccl",
+            pg_options=FakeOptions(hccl_config={"hccl_buffer_size": 200}),
+            reuse_domain="shared",
+            create_fn=create_fn,
+        )
+
+        assert first is second
+        assert create_fn.call_count == 1
+        assert registry._entries[key].refcount == 2
+
+        assert registry.release(key) is None
+        assert registry._entries[key].refcount == 1
+        assert registry.release(key) == first
+        assert key not in registry._entries
+        destroy_fn.assert_called_once_with(first)
+
+
+def test_acquire_duplicate_non_group_member_handle_is_not_destroyed():
+    sentinel = _MODULE._load_non_group_member_sentinel()
+    destroy_fn = MagicMock()
+    registry = HcclPgRegistry()
+    key = make_hccl_pg_key([0, 1], "hccl", FakeOptions(), reuse_domain="shared")
+
+    existing_handle = MagicMock(name="existing_handle")
+
+    def create_fn():
+        registry._entries[key] = RegistryEntry(handle=existing_handle, refcount=1)
+        return sentinel
+
+    with _patch_destroy_process_group(destroy_fn), _set_non_group_member_sentinel(sentinel):
+        merged = registry.acquire(
+            ranks=[0, 1],
+            backend="hccl",
+            pg_options=FakeOptions(),
+            reuse_domain="shared",
+            create_fn=create_fn,
+        )
+
+    assert merged is existing_handle
+    assert registry._entries[key].refcount == 2
+    destroy_fn.assert_not_called()
+
+
+def test_clear_removes_entries_without_destroying_handles():
+    destroy_fn = MagicMock()
+    with _patch_destroy_process_group(destroy_fn):
+        registry = HcclPgRegistry()
+        key = make_hccl_pg_key(
+            [0, 1],
+            "hccl",
+            FakeOptions(hccl_config={"hccl_buffer_size": 200}),
+            reuse_domain="shared",
+        )
+        registry._entries[key] = RegistryEntry(handle=MagicMock(name="pg"), refcount=1)
+        registry.clear()
+
+    assert key not in registry._entries
+    destroy_fn.assert_not_called()
+
+
+def test_release_non_group_member_uses_actual_sentinel():
+    destroy_fn = MagicMock()
+    sentinel = _MODULE._load_non_group_member_sentinel()
+    with _patch_destroy_process_group(destroy_fn), _set_non_group_member_sentinel(sentinel):
+        registry = HcclPgRegistry()
+        key = make_hccl_pg_key([0, 1], "hccl", FakeOptions(), reuse_domain="shared")
+        registry._entries[key] = RegistryEntry(handle=sentinel, refcount=1)
+
+        assert registry.release(key) is None
+        assert key not in registry._entries
+        destroy_fn.assert_not_called()
+
+
+def test_acquire_fails_closed_when_unknown_non_default_option_is_present():
+    registry = HcclPgRegistry()
+    create_fn = MagicMock(side_effect=[MagicMock(name="first"), MagicMock(name="second")])
+
+    first = registry.acquire(
+        ranks=[0, 1],
+        backend="hccl",
+        pg_options=FakeOptionsWithUnknown(non_default_option=7),
+        reuse_domain="shared",
+        create_fn=create_fn,
+    )
+    second = registry.acquire(
+        ranks=[0, 1],
+        backend="hccl",
+        pg_options=FakeOptionsWithUnknown(non_default_option=7),
+        reuse_domain="shared",
+        create_fn=create_fn,
+    )
+
+    assert first is not second
+    assert create_fn.call_count == 2
+    assert not registry._entries
+
+
+def test_acquire_fails_closed_for_unknown_mapping_fields():
+    registry = HcclPgRegistry()
+    create_fn = MagicMock(side_effect=[MagicMock(name="first"), MagicMock(name="second")])
+
+    options = {"hccl_config": {"hccl_buffer_size": 200}, "non_default_field": 7}
+
+    first = registry.acquire(
+        ranks=[0, 1],
+        backend="hccl",
+        pg_options=options,
+        reuse_domain="shared",
+        create_fn=create_fn,
+    )
+    second = registry.acquire(
+        ranks=[0, 1],
+        backend="hccl",
+        pg_options=options,
+        reuse_domain="shared",
+        create_fn=create_fn,
+    )
+
+    assert first is not second
+    assert create_fn.call_count == 2
+    assert not registry._entries

tests/ut/patch/worker/patch_common/test_patch_distributed.py

@@ -12,101 +12,672 @@
 # limitations under the License.
 # This file is a part of the vllm-ascend project.
 #
+from __future__ import annotations
 
-from unittest.mock import MagicMock, patch
+import sys
+import weakref
+from contextlib import contextmanager
+from dataclasses import dataclass
+from datetime import timedelta
+from importlib.util import module_from_spec, spec_from_file_location
+from pathlib import Path
+from types import ModuleType, SimpleNamespace
+from typing import Any
+from unittest.mock import MagicMock, call
 
-import torch
-from vllm.distributed.parallel_state import GroupCoordinator
+import pytest
 
-from tests.ut.base import TestBase
-from vllm_ascend.patch.worker.patch_common.patch_distributed import \
-    GroupCoordinatorPatch
+_WORKTREE_ROOT = Path(__file__).resolve().parents[5]
+_PATCH_MODULE_PATH = _WORKTREE_ROOT / "vllm_ascend/patch/worker/patch_distributed.py"
+_PATCH_MODULE_NAME = "vllm_ascend.patch.worker.patch_distributed"
+_REGISTRY_MODULE_PATH = _WORKTREE_ROOT / "vllm_ascend/patch/worker/_hccl_pg_registry.py"
+_REGISTRY_MODULE_NAME = "vllm_ascend.patch.worker._hccl_pg_registry"
 
 
-class TestPatchDistributed(TestBase):
+class FakeBackend(str):
+    pass
 
-    def setUp(self):
-        self.mock_group_ranks = [[0, 1]]
-        self.mock_local_rank = 0
-        self.mock_backend = "hccl"
-        self.mock_use_device_comm = True
 
-        patcher_get_rank = patch("torch.distributed.get_rank", return_value=0)
-        patcher_new_group = patch("torch.distributed.new_group",
-                                  return_value=MagicMock())
-        patcher_is_cuda_alike = patch(
-            "vllm.platforms.current_platform.is_cuda_alike", return_value=True)
-        patcher_device_comm_cls = patch(
-            "vllm.distributed.parallel_state.resolve_obj_by_qualname",
-            return_value=MagicMock())
+class FakeTensor:
+    def __init__(self, shape: tuple[int, ...]):
+        self._shape = shape
 
-        self.mock_get_rank = patcher_get_rank.start()
-        self.mock_new_group = patcher_new_group.start()
-        self.mock_is_cuda_alike = patcher_is_cuda_alike.start()
-        self.mock_resolve_obj = patcher_device_comm_cls.start()
+    def dim(self) -> int:
+        return len(self._shape)
 
-        self.addCleanup(patcher_get_rank.stop)
-        self.addCleanup(patcher_new_group.stop)
-        self.addCleanup(patcher_is_cuda_alike.stop)
-        self.addCleanup(patcher_device_comm_cls.stop)
+    def size(self) -> tuple[int, ...]:
+        return self._shape
 
-        self.group_coordinator = GroupCoordinatorPatch(
-            group_ranks=self.mock_group_ranks,
-            local_rank=self.mock_local_rank,
-            torch_distributed_backend=self.mock_backend,
-            use_device_communicator=self.mock_use_device_comm)
 
-    def test_GroupCoordinator_patched(self):
-        self.assertIs(GroupCoordinator, GroupCoordinatorPatch)
+class FakeProcessGroup:
+    def __init__(
+        self,
+        backend: str,
+        ranks: tuple[int, ...],
+        sequence: int,
+        pg_options: object | None = None,
+    ):
+        self.backend = backend
+        self.ranks = ranks
+        self.sequence = sequence
+        self.pg_options = pg_options
 
-    def test_all_to_all_returns_input_when_world_size_1(self):
-        self.group_coordinator.world_size = 1
-        input_tensor = torch.randn(2, 3)
-        output = self.group_coordinator.all_to_all(input_tensor)
-        self.assertTrue(torch.equal(output, input_tensor))
 
-    def test_all_to_all_raises_assertion_on_invalid_scatter_dim(self):
-        input_tensor = torch.randn(2, 3)
-        with self.assertRaises(AssertionError) as cm:
-            self.group_coordinator.all_to_all(input_tensor, scatter_dim=2)
-        self.assertIn("Invalid scatter dim", str(cm.exception))
+@dataclass
+class RealisticFakeHcclOptions:
+    backend: str = "hccl"
+    global_ranks_in_group: list[int] | tuple[int, ...] = ()
+    group_id: str = ""
+    group_name: str = ""
+    hccl_config: dict[str, int] | None = None
+    is_high_priority_stream: bool = False
+    op_timeout: object = timedelta(seconds=10)
 
-    def test_all_to_all_raises_assertion_on_invalid_gather_dim(self):
-        input_tensor = torch.randn(2, 3)
-        with self.assertRaises(AssertionError) as cm:
-            self.group_coordinator.all_to_all(input_tensor, gather_dim=2)
-        self.assertIn("Invalid gather dim", str(cm.exception))
 
-    def test_all_to_all_calls_device_communicator_with_correct_args(self):
-        mock_communicator = MagicMock()
-        self.group_coordinator.device_communicator = mock_communicator
+def _load_module(module_name: str, module_path: Path) -> Any:
+    spec = spec_from_file_location(module_name, str(module_path))
+    if spec is None or spec.loader is None:
+        raise RuntimeError(f"Failed to load module spec for {module_name}")
+    module = module_from_spec(spec)
+    sys.modules[module_name] = module
+    spec.loader.exec_module(module)
+    return module
 
-        input_tensor = torch.randn(2, 3)
-        scatter_dim = 0
-        gather_dim = 1
-        scatter_sizes = [1, 1]
-        gather_sizes = [1, 1]
 
-        self.group_coordinator.all_to_all(input_tensor,
-                                          scatter_dim=scatter_dim,
-                                          gather_dim=gather_dim,
-                                          scatter_sizes=scatter_sizes,
-                                          gather_sizes=gather_sizes)
+@contextmanager
+def _load_patch_distributed_module():
+    non_group_member = object()
+    new_group_calls: list[dict[str, object]] = []
+    destroy_process_group = MagicMock()
+    destroy_distributed_environment = MagicMock(name="destroy_distributed_environment")
+    get_rank = MagicMock(return_value=0)
+    current_device = MagicMock(return_value="npu:0")
+    communicator_instances: list[object] = []
+    unique_name_counter = {"value": 0}
+    sequence_counter = {"value": 0}
+    registered_groups = {}
+    shared_hccl_options = {"hccl_config": {"hccl_buffer_size": 200}}
 
-        mock_communicator.all_to_all.assert_called_once_with(
-            input_tensor, scatter_dim, gather_dim, scatter_sizes, gather_sizes)
+    torch_module: Any = ModuleType("torch")
+    torch_distributed: Any = ModuleType("torch.distributed")
+    torch_distributed_c10d: Any = ModuleType("torch.distributed.distributed_c10d")
 
-    def test_all_to_all_calls_device_communicator_without_sizes(self):
-        mock_communicator = MagicMock()
-        self.group_coordinator.device_communicator = mock_communicator
+    def new_group(ranks, backend, pg_options=None):
+        backend_name = str(backend)
+        new_group_calls.append(
+            {
+                "ranks": tuple(ranks),
+                "backend": backend_name,
+                "pg_options": pg_options,
+            }
+        )
+        if get_rank() not in ranks:
+            return non_group_member
+        handle = FakeProcessGroup(
+            backend=backend_name,
+            ranks=tuple(ranks),
+            sequence=sequence_counter["value"],
+            pg_options=pg_options,
+        )
+        sequence_counter["value"] += 1
+        return handle
 
-        input_tensor = torch.randn(2, 3)
-        scatter_dim = 0
-        gather_dim = 1
+    class GroupMember:
+        NON_GROUP_MEMBER = non_group_member
 
-        self.group_coordinator.all_to_all(input_tensor,
-                                          scatter_dim=scatter_dim,
-                                          gather_dim=gather_dim)
+    torch_distributed.Backend = FakeBackend
+    torch_distributed.get_rank = get_rank
+    torch_distributed.new_group = new_group
+    torch_distributed.destroy_process_group = destroy_process_group
+    torch_distributed.distributed_c10d = torch_distributed_c10d
+    torch_distributed_c10d.GroupMember = GroupMember
 
-        mock_communicator.all_to_all.assert_called_once_with(
-            input_tensor, scatter_dim, gather_dim, None, None)
+    torch_module.Tensor = FakeTensor
+    torch_module.distributed = torch_distributed
+    torch_module.equal = lambda lhs, rhs: lhs is rhs
+    torch_module.randn = lambda *shape: FakeTensor(shape)
+    torch_module.npu = SimpleNamespace(current_device=current_device)
+
+    vllm_module: Any = ModuleType("vllm")
+    vllm_distributed: Any = ModuleType("vllm.distributed")
+    parallel_state_module: Any = ModuleType("vllm.distributed.parallel_state")
+
+    class BaseGroupCoordinator:
+        pass
+
+    def _get_unique_name(group_name: str) -> str:
+        unique_name_counter["value"] += 1
+        return f"{group_name}-{unique_name_counter['value']}"
+
+    def _register_group(group):
+        registered_groups[group.unique_name] = weakref.ref(group)
+
+    parallel_state_module.GroupCoordinator = BaseGroupCoordinator
+    parallel_state_module._get_unique_name = _get_unique_name
+    parallel_state_module._register_group = MagicMock(side_effect=_register_group)
+    parallel_state_module._groups = registered_groups
+    parallel_state_module.destroy_distributed_environment = destroy_distributed_environment
+
+    shm_broadcast_module: Any = ModuleType("vllm.distributed.device_communicators.shm_broadcast")
+
+    class MessageQueue:
+        create_from_process_group = MagicMock(side_effect=lambda group, *_: SimpleNamespace(group=group))
+
+    shm_broadcast_module.MessageQueue = MessageQueue
+
+    vllm_distributed.parallel_state = parallel_state_module
+    vllm_distributed.destroy_distributed_environment = destroy_distributed_environment
+    vllm_module.distributed = vllm_distributed
+
+    vllm_ascend_module: Any = ModuleType("vllm_ascend")
+    vllm_ascend_patch: Any = ModuleType("vllm_ascend.patch")
+    vllm_ascend_patch_worker: Any = ModuleType("vllm_ascend.patch.worker")
+    vllm_ascend_distributed: Any = ModuleType("vllm_ascend.distributed")
+    vllm_ascend_device_communicators: Any = ModuleType("vllm_ascend.distributed.device_communicators")
+    npu_communicator_module: Any = ModuleType("vllm_ascend.distributed.device_communicators.npu_communicator")
+    utils_module: Any = ModuleType("vllm_ascend.utils")
+
+    class FakeNPUCommunicator:
+        def __init__(self, **kwargs):
+            self.init_kwargs = kwargs
+            self.destroy = MagicMock()
+            self.all_to_all = MagicMock()
+            communicator_instances.append(self)
+
+    npu_communicator_module.NPUCommunicator = FakeNPUCommunicator
+    utils_module.create_hccl_pg_options = MagicMock(return_value=shared_hccl_options)
+
+    vllm_ascend_module.patch = vllm_ascend_patch
+    vllm_ascend_patch.worker = vllm_ascend_patch_worker
+    vllm_ascend_module.distributed = vllm_ascend_distributed
+    vllm_ascend_distributed.device_communicators = vllm_ascend_device_communicators
+
+    modules = {
+        "torch": torch_module,
+        "torch.distributed": torch_distributed,
+        "torch.distributed.distributed_c10d": torch_distributed_c10d,
+        "vllm": vllm_module,
+        "vllm.distributed": vllm_distributed,
+        "vllm.distributed.parallel_state": parallel_state_module,
+        "vllm.distributed.device_communicators.shm_broadcast": shm_broadcast_module,
+        "vllm_ascend": vllm_ascend_module,
+        "vllm_ascend.patch": vllm_ascend_patch,
+        "vllm_ascend.patch.worker": vllm_ascend_patch_worker,
+        "vllm_ascend.distributed": vllm_ascend_distributed,
+        "vllm_ascend.distributed.device_communicators": (vllm_ascend_device_communicators),
+        "vllm_ascend.distributed.device_communicators.npu_communicator": (npu_communicator_module),
+        "vllm_ascend.utils": utils_module,
+    }
+
+    previous_modules = {name: sys.modules.get(name) for name in modules}
+    previous_patch_module = sys.modules.get(_PATCH_MODULE_NAME)
+    previous_registry_module = sys.modules.get(_REGISTRY_MODULE_NAME)
+    try:
+        sys.modules.update(modules)
+        sys.modules.pop(_PATCH_MODULE_NAME, None)
+        sys.modules.pop(_REGISTRY_MODULE_NAME, None)
+        registry_module = _load_module(_REGISTRY_MODULE_NAME, _REGISTRY_MODULE_PATH)
+        patch_module = _load_module(_PATCH_MODULE_NAME, _PATCH_MODULE_PATH)
+        yield SimpleNamespace(
+            module=patch_module,
+            registry_module=registry_module,
+            torch=torch_module,
+            distributed=torch_distributed,
+            parallel_state_module=parallel_state_module,
+            utils_module=utils_module,
+            new_group_calls=new_group_calls,
+            destroy_process_group=destroy_process_group,
+            destroy_distributed_environment=destroy_distributed_environment,
+            get_rank=get_rank,
+            current_device=current_device,
+            communicator_instances=communicator_instances,
+            non_group_member=non_group_member,
+            Backend=FakeBackend,
+            vllm_distributed=vllm_distributed,
+        )
+    finally:
+        if previous_patch_module is None:
+            sys.modules.pop(_PATCH_MODULE_NAME, None)
+        else:
+            sys.modules[_PATCH_MODULE_NAME] = previous_patch_module
+        if previous_registry_module is None:
+            sys.modules.pop(_REGISTRY_MODULE_NAME, None)
+        else:
+            sys.modules[_REGISTRY_MODULE_NAME] = previous_registry_module
+        for name, previous in previous_modules.items():
+            if previous is None:
+                sys.modules.pop(name, None)
+            else:
+                sys.modules[name] = previous
+
+
+@pytest.fixture
+def module_env():
+    with _load_patch_distributed_module() as env:
+        yield env
+
+
+def _make_group(
+    module_env,
+    *,
+    group_ranks: list[list[int]] | None = None,
+    local_rank: int = 0,
+    backend: str | FakeBackend = "hccl",
+    use_device_communicator: bool = False,
+    use_message_queue_broadcaster: bool = False,
+    group_name: str | None = None,
+):
+    return module_env.module.GroupCoordinatorPatch(
+        group_ranks=group_ranks or [[0, 1]],
+        local_rank=local_rank,
+        torch_distributed_backend=backend,
+        use_device_communicator=use_device_communicator,
+        use_message_queue_broadcaster=use_message_queue_broadcaster,
+        group_name=group_name,
+    )
+
+
+def _calls_with_backend(module_env, backend: str) -> list[dict[str, object]]:
+    return [call_entry for call_entry in module_env.new_group_calls if call_entry["backend"] == backend]
+
+
+def test_group_coordinator_is_patched(module_env):
+    assert module_env.parallel_state_module.GroupCoordinator is module_env.module.GroupCoordinatorPatch
+
+
+def test_same_hccl_group_reuses_device_pg_once(module_env):
+    first = _make_group(
+        module_env,
+        backend=module_env.Backend("hccl"),
+        group_name="tp",
+    )
+    second = _make_group(module_env, backend="hccl", group_name="world")
+
+    hccl_calls = _calls_with_backend(module_env, "hccl")
+    gloo_calls = _calls_with_backend(module_env, "gloo")
+
+    assert len(hccl_calls) == 1
+    assert len(gloo_calls) == 2
+    assert first.device_group is second.device_group
+
+
+def test_same_hccl_group_reuses_with_realistic_options_object(module_env):
+    module_env.utils_module.create_hccl_pg_options.return_value = RealisticFakeHcclOptions(
+        hccl_config={"hccl_buffer_size": 200}
+    )
+
+    first = _make_group(module_env, backend="hccl", group_name="tp")
+    second = _make_group(module_env, backend="hccl", group_name="world")
+
+    hccl_calls = _calls_with_backend(module_env, "hccl")
+    gloo_calls = _calls_with_backend(module_env, "gloo")
+
+    assert len(hccl_calls) == 1
+    assert len(gloo_calls) == 2
+    assert first.device_group is second.device_group
+
+
+def test_eplb_stays_isolated_from_ep_even_when_pg_options_match(module_env):
+    first = _make_group(module_env, group_name="ep")
+    second = _make_group(module_env, group_name="eplb")
+
+    hccl_calls = _calls_with_backend(module_env, "hccl")
+
+    assert len(hccl_calls) == 2
+    assert first.device_group is not second.device_group
+
+
+def test_mc2_stays_isolated_from_ep_even_when_pg_options_match(module_env):
+    first = _make_group(module_env, group_name="ep")
+    second = _make_group(module_env, group_name="mc2")
+
+    hccl_calls = _calls_with_backend(module_env, "hccl")
+
+    assert len(hccl_calls) == 2
+    assert first.device_group is not second.device_group
+
+
+def test_dynamic_eplb_stays_separate_from_ep_when_pg_options_differ(module_env):
+    default_hccl_pg_options = module_env.utils_module.create_hccl_pg_options.return_value
+
+    def fake_create_hccl_pg_options(group_name: str):
+        if group_name == "dynamic_eplb":
+            return {"hccl_config": {"hccl_buffer_size": 512}}
+        return default_hccl_pg_options
+
+    module_env.utils_module.create_hccl_pg_options.side_effect = fake_create_hccl_pg_options
+
+    first = _make_group(module_env, group_name="ep")
+    second = _make_group(module_env, group_name="dynamic_eplb")
+
+    hccl_calls = _calls_with_backend(module_env, "hccl")
+
+    assert len(hccl_calls) == 2
+    assert first.device_group is not second.device_group
+
+
+def test_unknown_groups_share_by_default_when_ranks_and_options_match(module_env):
+    first = _make_group(module_env, group_name="fc3_quant_x")
+    second = _make_group(module_env, group_name="fc3_quant_y")
+
+    hccl_calls = _calls_with_backend(module_env, "hccl")
+
+    assert module_env.module._resolve_reuse_domain("fc3_quant_x:0") == "shared"
+    assert len(hccl_calls) == 1
+    assert first.device_group is second.device_group
+
+
+def test_hccl_pg_options_are_recreated_for_each_group_ranks_entry(module_env):
+    _make_group(
+        module_env,
+        group_ranks=[[0], [1]],
+        group_name="tp",
+    )
+
+    assert module_env.utils_module.create_hccl_pg_options.call_count == 2
+
+
+def test_destroy_releases_all_acquired_keys_in_reverse_order(module_env):
+    group = _make_group(
+        module_env,
+        group_ranks=[[0, 1], [2, 3]],
+        group_name="tp",
+        use_device_communicator=True,
+        use_message_queue_broadcaster=True,
+    )
+    release_mock = MagicMock(wraps=module_env.module._HCCL_PG_REGISTRY.release)
+    module_env.module._HCCL_PG_REGISTRY.release = release_mock
+
+    cpu_group = group.cpu_group
+    shared_device_group = group.device_group
+    communicator = group.device_communicator
+    acquired_keys = list(group._acquired_hccl_keys)
+
+    destroy_order = []
+    communicator.destroy.side_effect = lambda: destroy_order.append("communicator")
+    module_env.destroy_process_group.side_effect = lambda group: destroy_order.append(group)
+
+    group.destroy()
+    group.destroy()
+
+    assert len(acquired_keys) == 2
+    assert release_mock.call_args_list == [call(acquired_keys[1]), call(acquired_keys[0])]
+    assert module_env.destroy_process_group.call_args_list == [call(shared_device_group), call(cpu_group)]
+    assert destroy_order == ["communicator", shared_device_group, cpu_group]
+    assert group.device_communicator is None
+    assert group.mq_broadcaster is None
+    assert not hasattr(group, "cpu_group")
+    assert not hasattr(group, "device_group")
+    assert group._acquired_hccl_keys == []
+
+
+def test_failed_cpu_group_init_rolls_back_acquired_hccl_keys(module_env):
+    original_new_group = module_env.distributed.new_group
+    release_mock = MagicMock(wraps=module_env.module._HCCL_PG_REGISTRY.release)
+    module_env.module._HCCL_PG_REGISTRY.release = release_mock
+
+    def failing_new_group(ranks, backend, pg_options=None):
+        if str(backend) == "gloo":
+            raise RuntimeError("gloo failed")
+        return original_new_group(ranks, backend, pg_options)
+
+    module_env.distributed.new_group = failing_new_group
+    hccl_key = module_env.registry_module.make_hccl_pg_key(
+        [0, 1],
+        "hccl",
+        module_env.utils_module.create_hccl_pg_options.return_value,
+        reuse_domain="shared",
+    )
+
+    with pytest.raises(RuntimeError, match="gloo failed"):
+        _make_group(
+            module_env,
+            group_ranks=[[0, 1]],
+            group_name="tp",
+        )
+
+    assert release_mock.call_args_list == [call(hccl_key)]
+    assert module_env.module._HCCL_PG_REGISTRY._entries == {}
+
+
+def test_failed_device_communicator_init_releases_all_keys_in_reverse_order(
+    module_env,
+):
+    release_mock = MagicMock(wraps=module_env.module._HCCL_PG_REGISTRY.release)
+    module_env.module._HCCL_PG_REGISTRY.release = release_mock
+    module_env.module.NPUCommunicator = MagicMock(side_effect=RuntimeError("communicator failed"))
+
+    key_a = module_env.registry_module.make_hccl_pg_key(
+        [0, 1],
+        "hccl",
+        module_env.utils_module.create_hccl_pg_options.return_value,
+        reuse_domain="shared",
+    )
+    key_b = module_env.registry_module.make_hccl_pg_key(
+        [2, 3],
+        "hccl",
+        module_env.utils_module.create_hccl_pg_options.return_value,
+        reuse_domain="shared",
+    )
+
+    with pytest.raises(RuntimeError, match="communicator failed"):
+        _make_group(
+            module_env,
+            group_ranks=[[0, 1], [2, 3]],
+            group_name="tp",
+            use_device_communicator=True,
+        )
+
+    assert release_mock.call_args_list == [call(key_b), call(key_a)]
+    assert module_env.module._HCCL_PG_REGISTRY._entries == {}
+
+
+def test_shared_hccl_group_is_destroyed_only_after_last_coordinator(module_env):
+    first = _make_group(
+        module_env,
+        group_ranks=[[0, 1]],
+        group_name="tp",
+    )
+    second = _make_group(
+        module_env,
+        group_ranks=[[0, 1]],
+        group_name="world",
+    )
+
+    cpu_group_first = first.cpu_group
+    cpu_group_second = second.cpu_group
+    shared_device_group = first.device_group
+
+    assert shared_device_group is second.device_group
+
+    first.destroy()
+
+    assert module_env.destroy_process_group.call_args_list == [call(cpu_group_first)]
+
+    second.destroy()
+
+    assert module_env.destroy_process_group.call_args_list == [
+        call(cpu_group_first),
+        call(shared_device_group),
+        call(cpu_group_second),
+    ]
+
+
+def test_destroy_distributed_environment_clears_registry_before_reinit(module_env):
+    group = _make_group(
+        module_env,
+        group_ranks=[[0, 1]],
+        group_name="tp",
+    )
+    first_device_group = group.device_group
+    call_observations: list[int] = []
+
+    def record_destroy():
+        call_observations.append(len(module_env.module._HCCL_PG_REGISTRY._entries))
+        return "destroyed"
+
+    module_env.destroy_distributed_environment.side_effect = record_destroy
+
+    assert len(_calls_with_backend(module_env, "hccl")) == 1
+    assert (
+        module_env.parallel_state_module.destroy_distributed_environment
+        is module_env.vllm_distributed.destroy_distributed_environment
+    )
+
+    result = module_env.parallel_state_module.destroy_distributed_environment()
+
+    assert result == "destroyed"
+    assert call_observations == [1]
+    assert module_env.module._HCCL_PG_REGISTRY._entries == {}
+
+    second_group = _make_group(
+        module_env,
+        group_ranks=[[0, 1]],
+        group_name="tp",
+    )
+
+    assert len(_calls_with_backend(module_env, "hccl")) == 2
+    assert second_group.device_group is not first_device_group
+
+
+def test_destroy_cleans_up_fail_closed_hccl_device_group(module_env):
+    module_env.utils_module.create_hccl_pg_options.return_value = {
+        "hccl_config": {"hccl_buffer_size": 200},
+        "non_default_field": 7,
+    }
+    group = _make_group(
+        module_env,
+        group_ranks=[[0, 1], [2, 3]],
+        group_name="tp",
+    )
+
+    cpu_group = group.cpu_group
+    device_group = group.device_group
+
+    assert group._acquired_hccl_keys == []
+
+    group.destroy()
+    group.destroy()
+
+    assert module_env.destroy_process_group.call_args_list == [
+        call(device_group),
+        call(cpu_group),
+    ]
+    assert group._acquired_hccl_keys == []
+    assert not hasattr(group, "cpu_group")
+    assert not hasattr(group, "device_group")
+
+
+def test_hccl_sleep_destroy_and_restore_shared_group(module_env):
+    group = _make_group(
+        module_env,
+        group_ranks=[[0, 1]],
+        group_name="tp",
+        use_device_communicator=True,
+    )
+    original_device_group = group.device_group
+    original_communicator = group.device_communicator
+
+    assert len(_calls_with_backend(module_env, "hccl")) == 1
+
+    assert group.destroy_hccl() is True
+
+    original_communicator.destroy.assert_called_once()
+    assert group.device_communicator is None
+    assert group.device_group is None
+    assert group._acquired_hccl_keys == []
+    assert module_env.destroy_process_group.call_args_list == [call(original_device_group)]
+
+    assert group.restore_hccl() is True
+
+    assert len(_calls_with_backend(module_env, "hccl")) == 2
+    assert group.device_group is not None
+    assert group.device_group is not original_device_group
+    assert group.device_communicator is not None
+    assert group.device_communicator is not original_communicator
+    assert group.device == "npu:0"
+
+    assert group.restore_hccl() is False
+
+
+def test_non_hccl_destroy_path_destroys_device_group_directly(module_env):
+    group = _make_group(
+        module_env,
+        backend="nccl",
+        group_name="tp",
+    )
+
+    cpu_group = group.cpu_group
+    device_group = group.device_group
+
+    group.destroy()
+    group.destroy()
+
+    assert module_env.destroy_process_group.call_args_list == [
+        call(device_group),
+        call(cpu_group),
+    ]
+    assert not hasattr(group, "cpu_group")
+    assert not hasattr(group, "device_group")
+
+
+def test_all_to_all_returns_input_when_world_size_is_one(module_env):
+    group = _make_group(module_env)
+    group.world_size = 1
+    input_tensor = module_env.torch.randn(2, 3)
+
+    assert group.all_to_all(input_tensor) is input_tensor
+
+
+def test_all_to_all_raises_assertion_on_invalid_scatter_dim(module_env):
+    group = _make_group(module_env)
+    input_tensor = module_env.torch.randn(2, 3)
+
+    with pytest.raises(AssertionError, match="Invalid scatter dim"):
+        group.all_to_all(input_tensor, scatter_dim=2)
+
+
+def test_all_to_all_raises_assertion_on_invalid_gather_dim(module_env):
+    group = _make_group(module_env)
+    input_tensor = module_env.torch.randn(2, 3)
+
+    with pytest.raises(AssertionError, match="Invalid gather dim"):
+        group.all_to_all(input_tensor, gather_dim=2)
+
+
+def test_all_to_all_calls_device_communicator_with_correct_args(module_env):
+    group = _make_group(module_env)
+    communicator = MagicMock()
+    communicator.all_to_all.return_value = "ok"
+    group.device_communicator = communicator
+
+    input_tensor = module_env.torch.randn(2, 3)
+    output = group.all_to_all(
+        input_tensor,
+        scatter_dim=0,
+        gather_dim=1,
+        scatter_sizes=[1, 1],
+        gather_sizes=[1, 1],
+    )
+
+    communicator.all_to_all.assert_called_once_with(
+        input_tensor,
+        0,
+        1,
+        [1, 1],
+        [1, 1],
+    )
+    assert output == "ok"
+
+
+def test_all_to_all_calls_device_communicator_without_sizes(module_env):
+    group = _make_group(module_env)
+    communicator = MagicMock()
+    communicator.all_to_all.return_value = "ok"
+    group.device_communicator = communicator
+
+    input_tensor = module_env.torch.randn(2, 3)
+    output = group.all_to_all(input_tensor, scatter_dim=0, gather_dim=1)
+
+    communicator.all_to_all.assert_called_once_with(input_tensor, 0, 1, None, None)
+    assert output == "ok"

tests/ut/patch/worker/test_patch_deepseek_v2.py

@@ -0,0 +1,36 @@
+# SPDX-License-Identifier: Apache-2.0
+
+from types import SimpleNamespace
+
+from vllm_ascend.patch.worker.patch_deepseek_v2 import _should_skip_indexer_init
+
+
+def _config(**overrides) -> SimpleNamespace:
+    values = {"num_hidden_layers": 80}
+    values.update(overrides)
+    return SimpleNamespace(**values)
+
+
+def test_glm51_skip_topk_keeps_per_layer_indexer():
+    assert not _should_skip_indexer_init(
+        _config(),
+        "model.layers.2.self_attn",
+        skip_topk=True,
+    )
+
+
+def test_glm52_shared_layer_skips_indexer_init():
+    assert _should_skip_indexer_init(
+        _config(indexer_types=["full", "full", "shared"]),
+        "model.layers.2.self_attn",
+        skip_topk=True,
+    )
+
+
+def test_mtp_layer_keeps_indexer():
+    indexer_types = ["full"] * 80 + ["shared"]
+    assert not _should_skip_indexer_init(
+        _config(indexer_types=indexer_types),
+        "model.layers.80.self_attn",
+        skip_topk=True,
+    )

tests/ut/patch/worker/test_patch_eagle3_pp_aux.py

@@ -0,0 +1,292 @@
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# This file is a part of the vllm-ascend project.
+#
+
+from types import SimpleNamespace
+
+import torch
+import torch.nn as nn
+from vllm.model_executor.models.interfaces import EagleModelMixin
+from vllm.sequence import IntermediateTensors
+
+from vllm_ascend.patch.worker import patch_eagle3_pp_aux as eagle3_pp_aux
+
+
+class _FakePPGroup:
+    def __init__(self, is_first_rank: bool, is_last_rank: bool):
+        self.is_first_rank = is_first_rank
+        self.is_last_rank = is_last_rank
+
+
+class _FakeLayer(nn.Module):
+    def __init__(self, delta: float):
+        super().__init__()
+        self.delta = delta
+
+    def forward(self, positions, hidden_states, residual, kv_cache, attn_metadata, llama_4_scaling):
+        del positions, kv_cache, attn_metadata, llama_4_scaling
+        next_hidden_states = hidden_states + self.delta
+        next_residual = torch.zeros_like(hidden_states) if residual is None else residual + self.delta
+        return next_hidden_states, next_residual
+
+
+class _FakeEagleMixinLayer(nn.Module):
+    def __init__(self, delta: float):
+        super().__init__()
+        self.delta = delta
+
+    def forward(self, positions, hidden_states, residual):
+        del positions
+        next_hidden_states = hidden_states + self.delta
+        next_residual = torch.zeros_like(hidden_states) if residual is None else residual + self.delta
+        return next_hidden_states, next_residual
+
+
+class _FakeDeepseekV2Model(nn.Module):
+    def __init__(self, start_layer: int, end_layer: int, aux_hidden_state_layers: tuple[int, ...]):
+        super().__init__()
+        self.start_layer = start_layer
+        self.end_layer = end_layer
+        self.aux_hidden_state_layers = aux_hidden_state_layers
+        self.config = SimpleNamespace(hidden_size=2)
+        self.layers = nn.ModuleList([_FakeLayer(float(i + 1)) for i in range(4)])
+
+    def embed_input_ids(self, input_ids):
+        return input_ids.to(torch.float32).unsqueeze(-1).expand(-1, self.config.hidden_size)
+
+    def norm(self, hidden_states, residual):
+        return hidden_states + residual, None
+
+    def make_empty_intermediate_tensors(self, batch_size, dtype, device):
+        return IntermediateTensors(
+            {
+                "hidden_states": torch.zeros((batch_size, self.config.hidden_size), dtype=dtype, device=device),
+                "residual": torch.zeros((batch_size, self.config.hidden_size), dtype=dtype, device=device),
+            }
+        )
+
+
+class _FakeEagleMixinModel(nn.Module, EagleModelMixin):
+    def __init__(self, start_layer: int, end_layer: int, aux_hidden_state_layers: tuple[int, ...]):
+        super().__init__()
+        self.start_layer = start_layer
+        self.end_layer = end_layer
+        self.aux_hidden_state_layers = aux_hidden_state_layers
+        self.config = SimpleNamespace(hidden_size=2)
+        self.layers = nn.ModuleList([_FakeEagleMixinLayer(float(i + 1)) for i in range(4)])
+
+    def embed_input_ids(self, input_ids):
+        return input_ids.to(torch.float32).unsqueeze(-1).expand(-1, self.config.hidden_size)
+
+    def norm(self, hidden_states, residual):
+        return hidden_states + residual, None
+
+    def make_empty_intermediate_tensors(self, batch_size, dtype, device):
+        return IntermediateTensors(
+            {
+                "hidden_states": torch.zeros((batch_size, self.config.hidden_size), dtype=dtype, device=device),
+                "residual": torch.zeros((batch_size, self.config.hidden_size), dtype=dtype, device=device),
+            }
+        )
+
+
+def test_extract_aux_from_intermediate_sorts_by_aux_index():
+    aux_2 = torch.full((1, 2), 2.0)
+    aux_10 = torch.full((1, 2), 10.0)
+    intermediate = IntermediateTensors(
+        {
+            "hidden_states": torch.zeros((1, 2)),
+            "aux_layer_10": aux_10,
+            "aux_layer_2": aux_2,
+        }
+    )
+
+    aux_states = eagle3_pp_aux._extract_aux_from_intermediate(intermediate)
+
+    assert len(aux_states) == 2
+    torch.testing.assert_close(aux_states[0], aux_2)
+    torch.testing.assert_close(aux_states[1], aux_10)
+
+
+def test_non_last_pp_rank_carries_previous_and_local_aux_states(monkeypatch):
+    monkeypatch.setattr(
+        eagle3_pp_aux,
+        "get_pp_group",
+        lambda: _FakePPGroup(is_first_rank=False, is_last_rank=False),
+    )
+    forward = eagle3_pp_aux._make_deepseek_v2_forward()
+    model = _FakeDeepseekV2Model(start_layer=2, end_layer=3, aux_hidden_state_layers=(1, 2))
+    previous_aux = torch.full((2, 2), 11.0)
+    hidden_states = torch.full((2, 2), 3.0)
+    residual = torch.full((2, 2), 5.0)
+    intermediate = IntermediateTensors(
+        {
+            "hidden_states": hidden_states,
+            "residual": residual,
+            "aux_layer_0": previous_aux,
+        }
+    )
+
+    output = forward(
+        model,
+        None,
+        torch.arange(2),
+        kv_caches=[None] * 4,
+        attn_metadata=None,
+        intermediate_tensors=intermediate,
+    )
+
+    assert isinstance(output, IntermediateTensors)
+    assert set(output.tensors) == {"hidden_states", "residual", "aux_layer_0", "aux_layer_1"}
+    torch.testing.assert_close(output["aux_layer_0"], previous_aux)
+    torch.testing.assert_close(output["aux_layer_1"], hidden_states + residual)
+    torch.testing.assert_close(output["hidden_states"], hidden_states + 3.0)
+    torch.testing.assert_close(output["residual"], residual + 3.0)
+
+
+def test_last_pp_rank_returns_complete_aux_states(monkeypatch):
+    monkeypatch.setattr(
+        eagle3_pp_aux,
+        "get_pp_group",
+        lambda: _FakePPGroup(is_first_rank=False, is_last_rank=True),
+    )
+    forward = eagle3_pp_aux._make_deepseek_v2_forward()
+    model = _FakeDeepseekV2Model(start_layer=3, end_layer=4, aux_hidden_state_layers=(1, 3))
+    previous_aux = torch.full((2, 2), 13.0)
+    hidden_states = torch.full((2, 2), 7.0)
+    residual = torch.full((2, 2), 2.0)
+    intermediate = IntermediateTensors(
+        {
+            "hidden_states": hidden_states,
+            "residual": residual,
+            "aux_layer_0": previous_aux,
+        }
+    )
+
+    output_hidden_states, aux_states = forward(
+        model,
+        None,
+        torch.arange(2),
+        kv_caches=[None] * 4,
+        attn_metadata=None,
+        intermediate_tensors=intermediate,
+    )
+
+    assert len(aux_states) == 2
+    torch.testing.assert_close(aux_states[0], previous_aux)
+    torch.testing.assert_close(aux_states[1], hidden_states + residual)
+    torch.testing.assert_close(output_hidden_states, hidden_states + 4.0 + residual + 4.0)
+
+
+def test_eagle_mixin_non_last_pp_rank_carries_previous_and_local_aux_states(monkeypatch):
+    monkeypatch.setattr(
+        eagle3_pp_aux,
+        "get_pp_group",
+        lambda: _FakePPGroup(is_first_rank=False, is_last_rank=False),
+    )
+    forward = eagle3_pp_aux._make_eagle_mixin_forward()
+    model = _FakeEagleMixinModel(start_layer=2, end_layer=3, aux_hidden_state_layers=(1, 3))
+    previous_aux = torch.full((2, 2), 11.0)
+    hidden_states = torch.full((2, 2), 3.0)
+    residual = torch.full((2, 2), 5.0)
+    intermediate = IntermediateTensors(
+        {
+            "hidden_states": hidden_states,
+            "residual": residual,
+            "aux_layer_0": previous_aux,
+        }
+    )
+
+    output = forward(
+        model,
+        None,
+        torch.arange(2),
+        intermediate_tensors=intermediate,
+    )
+
+    assert isinstance(output, IntermediateTensors)
+    assert set(output.tensors) == {"hidden_states", "residual", "aux_layer_0", "aux_layer_1"}
+    torch.testing.assert_close(output["aux_layer_0"], previous_aux)
+    torch.testing.assert_close(output["aux_layer_1"], hidden_states + 3.0 + residual + 3.0)
+    torch.testing.assert_close(output["hidden_states"], hidden_states + 3.0)
+    torch.testing.assert_close(output["residual"], residual + 3.0)
+
+
+def test_eagle_mixin_last_pp_rank_returns_complete_aux_states(monkeypatch):
+    monkeypatch.setattr(
+        eagle3_pp_aux,
+        "get_pp_group",
+        lambda: _FakePPGroup(is_first_rank=False, is_last_rank=True),
+    )
+    forward = eagle3_pp_aux._make_eagle_mixin_forward()
+    model = _FakeEagleMixinModel(start_layer=3, end_layer=4, aux_hidden_state_layers=(1, 4))
+    previous_aux = torch.full((2, 2), 13.0)
+    hidden_states = torch.full((2, 2), 7.0)
+    residual = torch.full((2, 2), 2.0)
+    intermediate = IntermediateTensors(
+        {
+            "hidden_states": hidden_states,
+            "residual": residual,
+            "aux_layer_0": previous_aux,
+        }
+    )
+
+    output_hidden_states, aux_states = forward(
+        model,
+        None,
+        torch.arange(2),
+        intermediate_tensors=intermediate,
+    )
+
+    assert len(aux_states) == 2
+    torch.testing.assert_close(aux_states[0], previous_aux)
+    torch.testing.assert_close(aux_states[1], hidden_states + 4.0 + residual + 4.0)
+    torch.testing.assert_close(output_hidden_states, hidden_states + 4.0 + residual + 4.0)
+
+
+def test_make_empty_intermediate_tensors_allocates_only_incoming_aux_layers():
+    model = _FakeDeepseekV2Model(start_layer=2, end_layer=4, aux_hidden_state_layers=(0, 2, 3))
+
+    eagle3_pp_aux._patch_make_empty_intermediate_tensors(model)
+    result = model.make_empty_intermediate_tensors(
+        batch_size=3,
+        dtype=torch.float32,
+        device=torch.device("cpu"),
+    )
+
+    assert set(result.tensors) == {"hidden_states", "residual", "aux_layer_0"}
+    assert result["aux_layer_0"].shape == (3, 2)
+    assert result["aux_layer_0"].dtype == torch.float32
+
+
+def test_patch_accepts_eagle_mixin_model():
+    model = _FakeEagleMixinModel(start_layer=2, end_layer=4, aux_hidden_state_layers=(0, 2, 3))
+
+    assert eagle3_pp_aux.patch_eagle3_pp_aux_propagation(model) is True
+    assert model._eagle3_pp_aux_forward_patched is True
+    assert model._eagle3_pp_aux_make_empty_patched is True
+
+    result = model.make_empty_intermediate_tensors(
+        batch_size=3,
+        dtype=torch.float32,
+        device=torch.device("cpu"),
+    )
+
+    assert set(result.tensors) == {"hidden_states", "residual", "aux_layer_0"}
+
+
+def test_patch_rejects_unsupported_model():
+    unsupported_model = nn.Linear(2, 2)
+
+    assert eagle3_pp_aux.patch_eagle3_pp_aux_propagation(unsupported_model) is False

tests/ut/patch/worker/test_patch_qwen3_5_mtp.py

@@ -0,0 +1,74 @@
+# SPDX-License-Identifier: Apache-2.0
+
+from types import SimpleNamespace
+from unittest.mock import MagicMock, patch
+
+import pytest
+import torch
+from vllm.sequence import IntermediateTensors
+
+from vllm_ascend.patch.worker import patch_qwen3_5
+
+
+@pytest.mark.skipif(
+    patch_qwen3_5.Qwen3_5MultiTokenPredictor is None,
+    reason="Qwen3.5 MTP model is not available in this vLLM version.",
+)
+def test_qwen3_5_mtp_forward_uses_local_inputs_on_last_pp_rank():
+    predictor = patch_qwen3_5.Qwen3_5MultiTokenPredictor.__new__(patch_qwen3_5.Qwen3_5MultiTokenPredictor)
+    predictor.num_mtp_layers = 2
+    predictor.embed_input_ids = MagicMock(return_value=torch.ones(2, 4))
+    predictor.pre_fc_norm_embedding = MagicMock(side_effect=lambda x: x + 1)
+    predictor.pre_fc_norm_hidden = MagicMock(side_effect=lambda x: x + 2)
+    predictor.fc = MagicMock(side_effect=lambda x: x[:, :4] + x[:, 4:])
+    layer0 = MagicMock(return_value=(torch.full((2, 4), 3.0), torch.full((2, 4), 4.0)))
+    layer1 = MagicMock(return_value=(torch.full((2, 4), 5.0), torch.full((2, 4), 6.0)))
+    predictor.layers = [layer0, layer1]
+    predictor.norm = MagicMock(return_value=(torch.full((2, 4), 7.0), None))
+
+    with patch(
+        "vllm_ascend.patch.worker.patch_qwen3_5.get_pp_group",
+        return_value=SimpleNamespace(is_last_rank=True),
+    ):
+        output = predictor.forward(
+            input_ids=torch.tensor([1, 2]),
+            positions=torch.tensor([0, 1]),
+            hidden_states=torch.zeros(2, 4),
+            intermediate_tensors=IntermediateTensors({"hidden_states": torch.full((2, 4), 99.0)}),
+            spec_step_idx=3,
+        )
+
+    predictor.embed_input_ids.assert_called_once()
+    layer1.assert_called_once()
+    predictor.norm.assert_called_once()
+    assert torch.equal(output, torch.full((2, 4), 7.0))
+
+
+@pytest.mark.skipif(
+    patch_qwen3_5.Qwen3_5MultiTokenPredictor is None,
+    reason="Qwen3.5 MTP model is not available in this vLLM version.",
+)
+def test_qwen3_5_mtp_forward_returns_intermediate_tensors_on_non_last_pp_rank():
+    predictor = patch_qwen3_5.Qwen3_5MultiTokenPredictor.__new__(patch_qwen3_5.Qwen3_5MultiTokenPredictor)
+    predictor.num_mtp_layers = 1
+    predictor.embed_input_ids = MagicMock(return_value=torch.ones(1, 4))
+    predictor.pre_fc_norm_embedding = MagicMock(side_effect=lambda x: x)
+    predictor.pre_fc_norm_hidden = MagicMock(side_effect=lambda x: x)
+    predictor.fc = MagicMock(side_effect=lambda x: x[:, :4])
+    predictor.layers = [MagicMock(return_value=(torch.full((1, 4), 3.0), torch.full((1, 4), 4.0)))]
+    predictor.norm = MagicMock()
+
+    with patch(
+        "vllm_ascend.patch.worker.patch_qwen3_5.get_pp_group",
+        return_value=SimpleNamespace(is_last_rank=False),
+    ):
+        output = predictor.forward(
+            input_ids=torch.tensor([1]),
+            positions=torch.tensor([0]),
+            hidden_states=torch.zeros(1, 4),
+        )
+
+    assert isinstance(output, IntermediateTensors)
+    assert torch.equal(output["hidden_states"], torch.full((1, 4), 3.0))
+    assert torch.equal(output["residual"], torch.full((1, 4), 4.0))
+    predictor.norm.assert_not_called()

tests/ut/patch/worker/test_patch_routed_experts_capture.py

@@ -0,0 +1,232 @@
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# This file is a part of the vllm-ascend project.
+#
+
+from __future__ import annotations
+
+from unittest.mock import MagicMock, patch
+
+import pytest
+import torch
+
+from vllm_ascend.patch.worker.patch_routed_experts_capture import capture
+
+
+class TestRoutedExpertsCapturerCapture:
+    """Unit tests for RoutedExpertsCapturer.capture method."""
+
+    def _create_mock_capturer(self, dp_rank=0, tp_size=1):
+        """Create a mock RoutedExpertsCapturer instance with required attributes."""
+        capturer = MagicMock()
+        capturer.dp_rank = dp_rank
+        capturer.tp_size = tp_size
+        capturer.device_buffer = torch.zeros((100, 10, 2), dtype=torch.int32)
+        return capturer
+
+    def _create_mock_forward_context(self, dp_metadata=None):
+        """Create a mock forward context."""
+        ctx = MagicMock()
+        ctx.dp_metadata = dp_metadata
+        return ctx
+
+    def _create_mock_dp_metadata(self, num_tokens_across_dp_cpu):
+        """Create mock DP metadata with token counts."""
+        dp_metadata = MagicMock()
+        dp_metadata.num_tokens_across_dp_cpu = torch.tensor(num_tokens_across_dp_cpu)
+        return dp_metadata
+
+    @patch("vllm_ascend.patch.worker.patch_routed_experts_capture.get_forward_context")
+    def test_single_dp(self, mock_get_ctx):
+        """Test single DP scenario (no data parallelism)."""
+        capturer = self._create_mock_capturer(dp_rank=0, tp_size=1)
+        ctx = self._create_mock_forward_context(dp_metadata=None)
+        mock_get_ctx.return_value = ctx
+
+        topk_ids = torch.tensor([[0, 1], [2, 3], [4, 5]], dtype=torch.int32)
+
+        capture(capturer, layer_id=0, topk_ids=topk_ids)
+
+        expected = topk_ids
+        actual = capturer.device_buffer[:3, 0, :]
+        torch.testing.assert_close(actual, expected)
+
+    @patch("vllm_ascend.patch.worker.patch_routed_experts_capture.get_forward_context")
+    def test_multi_dp_naive_dispatch(self, mock_get_ctx):
+        """Test multi-DP naive dispatch (n == total)."""
+        capturer = self._create_mock_capturer(dp_rank=0, tp_size=1)
+        dp_metadata = self._create_mock_dp_metadata([5, 7])
+        ctx = self._create_mock_forward_context(dp_metadata=dp_metadata)
+        mock_get_ctx.return_value = ctx
+
+        topk_ids = torch.arange(12 * 2).view(12, 2).to(torch.int32)
+
+        capture(capturer, layer_id=0, topk_ids=topk_ids)
+
+        expected = topk_ids[:5]
+        actual = capturer.device_buffer[:5, 0, :]
+        torch.testing.assert_close(actual, expected)
+
+    @patch("vllm_ascend.patch.worker.patch_routed_experts_capture.get_forward_context")
+    def test_multi_dp_modular_kernel(self, mock_get_ctx):
+        """Test multi-DP modular kernel path (n == token_num_per_dp)."""
+        capturer = self._create_mock_capturer(dp_rank=1, tp_size=1)
+        dp_metadata = self._create_mock_dp_metadata([5, 7])
+        ctx = self._create_mock_forward_context(dp_metadata=dp_metadata)
+        mock_get_ctx.return_value = ctx
+
+        topk_ids = torch.arange(7 * 2).view(7, 2).to(torch.int32)
+
+        capture(capturer, layer_id=0, topk_ids=topk_ids)
+
+        expected = topk_ids
+        actual = capturer.device_buffer[:7, 0, :]
+        torch.testing.assert_close(actual, expected)
+
+    @patch("vllm_ascend.patch.worker.patch_routed_experts_capture.get_forward_context")
+    def test_multi_dp_padded_all_gather(self, mock_get_ctx):
+        """Test multi-DP padded all-gather path (n == total_with_padding)."""
+        capturer = self._create_mock_capturer(dp_rank=0, tp_size=1)
+        dp_metadata = self._create_mock_dp_metadata([5, 7])
+        ctx = self._create_mock_forward_context(dp_metadata=dp_metadata)
+        mock_get_ctx.return_value = ctx
+
+        topk_ids = torch.arange(14 * 2).view(14, 2).to(torch.int32)
+
+        capture(capturer, layer_id=0, topk_ids=topk_ids)
+
+        expected = topk_ids[:5]
+        actual = capturer.device_buffer[:5, 0, :]
+        torch.testing.assert_close(actual, expected)
+
+    @patch("vllm_ascend.ascend_forward_context.get_forward_context")
+    @patch("vllm_ascend.patch.worker.patch_routed_experts_capture.get_forward_context")
+    @patch("vllm_ascend.patch.worker.patch_routed_experts_capture.get_tp_group")
+    @patch("vllm_ascend.patch.worker.patch_routed_experts_capture.dist")
+    def test_sp_modular_kernel_all2all(
+        self,
+        mock_dist,
+        mock_get_tp_group,
+        mock_get_ctx,
+        mock_get_ctx1,
+    ):
+        """Test SP + modular kernel path with ALLTOALL comm type."""
+        capturer = self._create_mock_capturer(dp_rank=0, tp_size=2)
+        dp_metadata = self._create_mock_dp_metadata([10])
+        ctx = self._create_mock_forward_context(dp_metadata=dp_metadata)
+        mock_get_ctx.return_value = ctx
+
+        mock_tp_group = MagicMock()
+        mock_get_tp_group.return_value = mock_tp_group
+        mock_tp_group.device_group = MagicMock()
+
+        from vllm_ascend.ascend_forward_context import _EXTRA_CTX, MoECommType
+
+        original_comm_type = getattr(_EXTRA_CTX, "moe_comm_type", None)
+        _EXTRA_CTX.moe_comm_type = MoECommType.ALLTOALL
+
+        try:
+            topk_ids = torch.arange(5 * 2).view(5, 2).to(torch.int32)
+
+            def mock_all_gather_impl(output_list, input_tensor, device_group):
+                output_list[0].copy_(input_tensor)
+                output_list[1].copy_(input_tensor + 5 * 2)
+
+            mock_dist.all_gather = mock_all_gather_impl
+
+            capture(capturer, layer_id=0, topk_ids=topk_ids)
+
+            expected = torch.arange(10 * 2).view(10, 2).to(torch.int32)
+            actual = capturer.device_buffer[:10, 0, :]
+            torch.testing.assert_close(actual, expected)
+        finally:
+            if original_comm_type is not None:
+                _EXTRA_CTX.moe_comm_type = original_comm_type
+            else:
+                delattr(_EXTRA_CTX, "moe_comm_type")
+
+    @patch("vllm_ascend.ascend_forward_context.get_forward_context")
+    @patch("vllm_ascend.patch.worker.patch_routed_experts_capture.get_forward_context")
+    @patch("vllm_ascend.patch.worker.patch_routed_experts_capture.get_tp_group")
+    @patch("vllm_ascend.patch.worker.patch_routed_experts_capture.dist")
+    def test_sp_modular_kernel_mc2(
+        self,
+        mock_dist,
+        mock_get_tp_group,
+        mock_get_ctx,
+        mock_get_ctx1,
+    ):
+        """Test SP + modular kernel path with MC2 comm type."""
+        capturer = self._create_mock_capturer(dp_rank=0, tp_size=2)
+        dp_metadata = self._create_mock_dp_metadata([5, 7])
+        ctx = self._create_mock_forward_context(dp_metadata=dp_metadata)
+        mock_get_ctx.return_value = ctx
+
+        mock_tp_group = MagicMock()
+        mock_get_tp_group.return_value = mock_tp_group
+        mock_tp_group.device_group = MagicMock()
+
+        from vllm_ascend.ascend_forward_context import _EXTRA_CTX, MoECommType
+
+        original_comm_type = getattr(_EXTRA_CTX, "moe_comm_type", None)
+        _EXTRA_CTX.moe_comm_type = MoECommType.MC2
+
+        try:
+            topk_ids = torch.arange(4 * 2).view(4, 2).to(torch.int32)
+
+            def mock_all_gather_impl(output_list, input_tensor, device_group):
+                output_list[0].copy_(input_tensor)
+                output_list[1].copy_(input_tensor + 4 * 2)
+
+            mock_dist.all_gather = mock_all_gather_impl
+
+            capture(capturer, layer_id=0, topk_ids=topk_ids)
+
+            actual = capturer.device_buffer[:5, 0, :]
+            assert actual.shape[0] == 5
+        finally:
+            if original_comm_type is not None:
+                _EXTRA_CTX.moe_comm_type = original_comm_type
+            else:
+                delattr(_EXTRA_CTX, "moe_comm_type")
+
+    @patch("vllm_ascend.patch.worker.patch_routed_experts_capture.get_forward_context")
+    def test_unexpected_batch_dim(self, mock_get_ctx):
+        """Test that unexpected batch dimension raises AssertionError."""
+        capturer = self._create_mock_capturer(dp_rank=0, tp_size=2)
+        dp_metadata = self._create_mock_dp_metadata([5, 7])
+        ctx = self._create_mock_forward_context(dp_metadata=dp_metadata)
+        mock_get_ctx.return_value = ctx
+
+        topk_ids = torch.randint(0, 8, (100, 2)).to(torch.int32)
+
+        with pytest.raises(AssertionError, match="unexpected topk_ids batch"):
+            capture(capturer, layer_id=0, topk_ids=topk_ids)
+
+    @patch("vllm_ascend.patch.worker.patch_routed_experts_capture.get_forward_context")
+    def test_layer_id_out_of_bounds(self, mock_get_ctx):
+        """Test that out-of-bounds layer_id is handled gracefully."""
+        capturer = self._create_mock_capturer(dp_rank=0, tp_size=1)
+        ctx = self._create_mock_forward_context(dp_metadata=None)
+        mock_get_ctx.return_value = ctx
+
+        topk_ids = torch.tensor([[0, 1]], dtype=torch.int32)
+        capturer.device_buffer = torch.zeros((100, 5, 2))
+
+        capture(capturer, layer_id=10, topk_ids=topk_ids)
+
+        assert torch.all(capturer.device_buffer == 0)
+
+
+if __name__ == "__main__":
+    pytest.main([__file__, "-v"])