Zhou TaichangandClaude Fable 5
10906e3a6d
Signed-off-by: tzhouam <tzhouam@connect.ust.hk> Co-authored-by: Claude Fable 5 <noreply@anthropic.com>
94 lines
3.6 KiB
Python
94 lines
3.6 KiB
Python
# SPDX-License-Identifier: Apache-2.0
|
|
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
|
"""Static coupling check for the ``has_preprocess`` capture/replay buffers.
|
|
|
|
For models with ``has_preprocess``, CUDA-graph *capture* (``_dummy_run``) and
|
|
*replay* (``_preprocess``) MUST read from the same pre-allocated buffers
|
|
(``self.input_ids.gpu`` + ``self.inputs_embeds.gpu``); otherwise replay reads a
|
|
buffer capture never wrote and produces garbage. This is a source-level check
|
|
(no GPU) that catches a one-sided edit to either branch.
|
|
|
|
The check is AST-scoped: it inspects ONLY the body of the ``has_preprocess``
|
|
guarded branch (not a line window), so a buffer read removed from that branch
|
|
cannot be masked by an adjacent ``else``/``elif`` that happens to read it. It
|
|
matches by attribute chain (``self.input_ids.gpu``), so local variable renames
|
|
do not break it.
|
|
|
|
Scope: the BASE runner (``OmniGPUModelRunner``) and the generation runner's
|
|
``_dummy_run`` override (which gained the branch in RFC #5450 C6; its
|
|
``_preprocess`` is inherited from the base).
|
|
"""
|
|
|
|
import ast
|
|
import inspect
|
|
import textwrap
|
|
|
|
import pytest
|
|
|
|
from vllm_omni.worker.gpu_generation_model_runner import GPUGenerationModelRunner
|
|
from vllm_omni.worker.gpu_model_runner import OmniGPUModelRunner
|
|
|
|
pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
|
|
|
|
_REQUIRED_BUFFERS = {"self.input_ids.gpu", "self.inputs_embeds.gpu"}
|
|
|
|
|
|
def _attr_chain(node: ast.AST) -> str | None:
|
|
"""Reconstruct a dotted ``self.x.y`` chain from an ast.Attribute node."""
|
|
parts: list[str] = []
|
|
cur = node
|
|
while isinstance(cur, ast.Attribute):
|
|
parts.append(cur.attr)
|
|
cur = cur.value
|
|
if isinstance(cur, ast.Name):
|
|
parts.append(cur.id)
|
|
return ".".join(reversed(parts))
|
|
return None
|
|
|
|
|
|
def _attr_chains_in(nodes: list[ast.stmt]) -> set[str]:
|
|
chains: set[str] = set()
|
|
for stmt in nodes:
|
|
for inner in ast.walk(stmt):
|
|
if isinstance(inner, ast.Attribute):
|
|
chain = _attr_chain(inner)
|
|
if chain is not None:
|
|
chains.add(chain)
|
|
return chains
|
|
|
|
|
|
def _has_preprocess_branch_reads_both_buffers(method) -> bool:
|
|
"""True iff SOME ``if has_preprocess:`` branch body reads BOTH buffers.
|
|
|
|
Only the guarded branch's own ``body`` is inspected (elif/else excluded),
|
|
which is what makes a one-sided buffer removal fail.
|
|
"""
|
|
src = textwrap.dedent(inspect.getsource(method))
|
|
tree = ast.parse(src)
|
|
guarded = [node for node in ast.walk(tree) if isinstance(node, ast.If) and "has_preprocess" in ast.dump(node.test)]
|
|
assert guarded, f"{method.__qualname__}: has_preprocess branch removed entirely"
|
|
return any(_REQUIRED_BUFFERS <= _attr_chains_in(node.body) for node in guarded)
|
|
|
|
|
|
@pytest.mark.parametrize(
|
|
("runner_cls", "method_name"),
|
|
[
|
|
(OmniGPUModelRunner, "_dummy_run"),
|
|
(OmniGPUModelRunner, "_preprocess"),
|
|
(GPUGenerationModelRunner, "_dummy_run"),
|
|
],
|
|
)
|
|
def test_has_preprocess_branch_reads_shared_buffers(runner_cls, method_name):
|
|
method = getattr(runner_cls, method_name)
|
|
assert _has_preprocess_branch_reads_both_buffers(method), (
|
|
f"{runner_cls.__name__}.{method_name}: the has_preprocess branch no "
|
|
"longer reads BOTH self.input_ids.gpu and self.inputs_embeds.gpu — capture "
|
|
"and replay would diverge."
|
|
)
|
|
|
|
|
|
def test_generation_dummy_run_overrides_with_branch():
|
|
# Guard against the override being deleted (which would silently fall back
|
|
# to the base implementation and vacuously pass the parametrized check).
|
|
assert "_dummy_run" in GPUGenerationModelRunner.__dict__
|