|
| 1 | +# SPDX-License-Identifier: Apache-2.0 |
| 2 | +"""Source-level regressions for Step3.5 MTP Ascend glue. |
| 3 | +
|
| 4 | +Importing the Step3.5 proposer can initialize runtime/device state in this |
| 5 | +branch. Keep these checks focused on cross-file contracts that are hard to |
| 6 | +exercise in a lightweight unit test, and avoid pinning the exact implementation |
| 7 | +sequence inside the proposer. |
| 8 | +""" |
| 9 | + |
| 10 | +from __future__ import annotations |
| 11 | + |
| 12 | +import ast |
| 13 | +from pathlib import Path |
| 14 | + |
| 15 | +ROOT = Path(__file__).resolve().parents[3] |
| 16 | +STEP3P5 = ROOT / "vllm_ascend" / "spec_decode" / "step3p5.py" |
| 17 | +BASE_PROPOSER = ROOT / "vllm_ascend" / "spec_decode" / "llm_base_proposer.py" |
| 18 | +PATCH_SPEC_CFG = ROOT / "vllm_ascend" / "patch" / "platform" / "patch_speculative_config.py" |
| 19 | +WORKER_PATCH_INIT = ROOT / "vllm_ascend" / "patch" / "worker" / "__init__.py" |
| 20 | +LEGACY_STEP3P7_PATCH = ROOT / "vllm_ascend" / "patch" / "worker" / "patch_step3p5_mtp.py" |
| 21 | + |
| 22 | + |
| 23 | +def _tree(path: Path) -> ast.Module: |
| 24 | + return ast.parse(path.read_text()) |
| 25 | + |
| 26 | + |
| 27 | +def _class(path: Path, name: str) -> ast.ClassDef: |
| 28 | + for node in _tree(path).body: |
| 29 | + if isinstance(node, ast.ClassDef) and node.name == name: |
| 30 | + return node |
| 31 | + raise AssertionError(f"class {name} not found in {path}") |
| 32 | + |
| 33 | + |
| 34 | +def _method(path: Path, cls_name: str, method_name: str) -> ast.FunctionDef: |
| 35 | + cls = _class(path, cls_name) |
| 36 | + for node in cls.body: |
| 37 | + if isinstance(node, ast.FunctionDef) and node.name == method_name: |
| 38 | + return node |
| 39 | + raise AssertionError(f"method {cls_name}.{method_name} not found") |
| 40 | + |
| 41 | + |
| 42 | +def _func(path: Path, name: str) -> ast.FunctionDef: |
| 43 | + for node in _tree(path).body: |
| 44 | + if isinstance(node, ast.FunctionDef) and node.name == name: |
| 45 | + return node |
| 46 | + raise AssertionError(f"function {name} not found in {path}") |
| 47 | + |
| 48 | + |
| 49 | +def _src(node: ast.AST) -> str: |
| 50 | + return ast.unparse(node) |
| 51 | + |
| 52 | + |
| 53 | +def test_step3p5_first_pass_forwards_rejected_token_counts() -> None: |
| 54 | + # set_inputs_first_pass is inherited from the base proposer; the base |
| 55 | + # simple-path is the canonical Step3.5 behaviour. Step3.5 only needs to |
| 56 | + # forward num_rejected_tokens_gpu through _propose. |
| 57 | + set_inputs = _method(BASE_PROPOSER, "AscendSpecDecodeBaseProposer", "set_inputs_first_pass") |
| 58 | + propose = _method(STEP3P5, "AscendStep3p5MTPProposer", "_propose") |
| 59 | + |
| 60 | + assert "num_rejected_tokens_gpu" in [arg.arg for arg in set_inputs.args.args] |
| 61 | + assert "num_rejected_tokens_gpu=num_rejected_tokens_gpu" in _src(propose) |
| 62 | + |
| 63 | + # Guard against the override creeping back: the previous step3p5 simple-path |
| 64 | + # was byte-equivalent to the base's `not needs_extra_input_slots and |
| 65 | + # pcp_size <= 1` branch, so a re-override is almost certainly redundant. |
| 66 | + step_methods = {n.name for n in _class(STEP3P5, "AscendStep3p5MTPProposer").body if isinstance(n, ast.FunctionDef)} |
| 67 | + assert "set_inputs_first_pass" not in step_methods |
| 68 | + |
| 69 | + |
| 70 | +def test_step3p5_draft_window_and_config_contracts() -> None: |
| 71 | + base_run = _method(BASE_PROPOSER, "AscendSpecDecodeBaseProposer", "_run_merged_draft") |
| 72 | + step_run = _method(STEP3P5, "AscendStep3p5MTPProposer", "_run_merged_draft") |
| 73 | + run_window = _src(_method(STEP3P5, "AscendStep3p5MTPProposer", "_run_window_draft_steps")) |
| 74 | + build_metadata = _src(_method(STEP3P5, "AscendStep3p5MTPProposer", "_build_step_attn_metadatas")) |
| 75 | + roll_inputs = _src(_method(STEP3P5, "AscendStep3p5MTPProposer", "_roll_window_inputs_only")) |
| 76 | + ensure_layer_types = _src( |
| 77 | + _method( |
| 78 | + STEP3P5, |
| 79 | + "AscendStep3p5MTPProposer", |
| 80 | + "_ensure_draft_layer_types_cover_mtp_layers", |
| 81 | + ) |
| 82 | + ) |
| 83 | + create_config = _src(_method(STEP3P5, "AscendStep3p5MTPProposer", "_create_draft_vllm_config")) |
| 84 | + |
| 85 | + assert [arg.arg for arg in step_run.args.args] == [arg.arg for arg in base_run.args.args] |
| 86 | + assert "multi_steps_attn_metadata.append(per_step_attn_metadata)" in build_metadata |
| 87 | + assert "multi_steps_attn_metadata[spec_step_idx]" in run_window |
| 88 | + assert "self.input_ids[token_indices_to_sample]" in roll_inputs |
| 89 | + assert "_ensure_draft_layer_types_cover_mtp_layers()" in create_config |
| 90 | + assert "self.draft_model_config.hf_config" in ensure_layer_types |
| 91 | + assert "self.vllm_config.model_config.hf_config" not in ensure_layer_types |
| 92 | + assert "sliding_attention" in ensure_layer_types |
| 93 | + |
| 94 | + |
| 95 | +def test_step3p7_uses_step3p5_mtp_override_without_legacy_runtime_patch() -> None: |
| 96 | + override_src = _src(_func(PATCH_SPEC_CFG, "hf_config_override")) |
| 97 | + |
| 98 | + assert "step3p7" in override_src |
| 99 | + assert "Step3p7ForConditionalGeneration" in override_src |
| 100 | + assert "step3p5_mtp" in override_src |
| 101 | + assert "Step3p5MTP" in override_src |
| 102 | + assert "patch_step3p5_mtp" not in WORKER_PATCH_INIT.read_text() |
| 103 | + assert not LEGACY_STEP3P7_PATCH.exists() |
0 commit comments