82 lines
4.0 KiB
Python
82 lines
4.0 KiB
Python
|
"""Idempotent source patches and CPU checks. No vLLM/torch imports or XPU needed."""
|
||
|
|
import argparse
|
||
|
|
import ast
|
||
|
|
import copy
|
||
|
|
import hashlib
|
||
|
|
import inspect
|
||
|
|
import json
|
||
|
|
from importlib import metadata
|
||
|
|
from pathlib import Path
|
||
|
|
from types import SimpleNamespace
|
||
|
|
|
||
|
|
from routing_compat import patch_source, regression_tests, INSERT
|
||
|
|
|
||
|
|
parser = argparse.ArgumentParser()
|
||
|
|
parser.add_argument('--site-packages', type=Path,
|
||
|
|
default=Path('/opt/vllm_kunlun/lib/python3.10/site-packages'))
|
||
|
|
parser.add_argument('--check', action='store_true', help='Verify without changing installed files')
|
||
|
|
args = parser.parse_args()
|
||
|
|
root = args.site_packages
|
||
|
|
route = root / 'vllm_kunlun/quantization/compressed_tensors/compressed_tensors_moe.py'
|
||
|
|
wrapper = root / 'vllm_kunlun/ops/_custom_ops.py'
|
||
|
|
lower = root / 'kunlun_ops/_xflashinfer_moe.py'
|
||
|
|
versions = {p: metadata.version(p) for p in ('vllm', 'vllm-kunlun', 'kunlun-ops')}
|
||
|
|
if not versions['vllm'].startswith('0.25.1') or not versions['vllm-kunlun'].startswith('0.25.1'):
|
||
|
|
raise RuntimeError(f'Unsupported environment: {versions}')
|
||
|
|
|
||
|
|
before_route = route.read_text()
|
||
|
|
after_route = patch_source(before_route)
|
||
|
|
# Reproduce the legacy bug and exercise patched method dispatch.
|
||
|
|
regression_tests(after_route.replace(INSERT, ''), after_route)
|
||
|
|
lower_fn = next(n for n in ast.parse(lower.read_text()).body
|
||
|
|
if isinstance(n, ast.FunctionDef) and n.name == 'moe_sigmoid_group_topk_norm')
|
||
|
|
signature = inspect.Signature([inspect.Parameter(a.arg, inspect.Parameter.POSITIONAL_OR_KEYWORD)
|
||
|
|
for a in lower_fn.args.args])
|
||
|
|
assert 'block_statistic' in signature.parameters
|
||
|
|
|
||
|
|
before_wrapper = wrapper.read_text()
|
||
|
|
lines = before_wrapper.splitlines(keepends=True)
|
||
|
|
names = {'moe_sigmoid_group_topk_norm', 'moe_sigmoid_group_topk_norm_cuda'}
|
||
|
|
functions = [n for n in ast.parse(before_wrapper).body if isinstance(n, ast.FunctionDef) and n.name in names]
|
||
|
|
assert len(functions) == 2
|
||
|
|
for fn in functions:
|
||
|
|
for call in ast.walk(fn):
|
||
|
|
if isinstance(call, ast.Call) and ast.unparse(call.func) == 'kunlun_ops.moe_sigmoid_group_topk_norm':
|
||
|
|
for kw in call.keywords:
|
||
|
|
if kw.arg == 'block_static':
|
||
|
|
idx = kw.lineno - 1
|
||
|
|
assert 'block_static=block_static' in lines[idx]
|
||
|
|
lines[idx] = lines[idx].replace('block_static=block_static', 'block_statistic=block_static')
|
||
|
|
after_wrapper = ''.join(lines)
|
||
|
|
for fn in [n for n in ast.parse(after_wrapper).body if isinstance(n, ast.FunctionDef) and n.name in names]:
|
||
|
|
node = copy.deepcopy(fn)
|
||
|
|
node.decorator_list = []; node.returns = None
|
||
|
|
for a in node.args.args:
|
||
|
|
a.annotation = None
|
||
|
|
calls = []
|
||
|
|
def spy(*a, **kw):
|
||
|
|
calls.append(signature.bind(*a, **kw).arguments)
|
||
|
|
scope = {'kunlun_ops': SimpleNamespace(moe_sigmoid_group_topk_norm=spy)}
|
||
|
|
exec(compile(ast.fix_missing_locations(ast.Module(body=[node], type_ignores=[])), '<wrapper-check>', 'exec'), scope)
|
||
|
|
values = {k: object() for k in ('x', 'topk_index', 'norm_score', 'block_static', 'bias')}
|
||
|
|
scope[fn.name](**values, scale=1.0, n_group=8, topk_group=4)
|
||
|
|
assert len(calls) == 1 and calls[0]['block_statistic'] is values['block_static']
|
||
|
|
assert calls[0]['bias'] is values['bias']
|
||
|
|
print('WRAPPER_PASS', fn.name)
|
||
|
|
|
||
|
|
for path, before, after in [(route, before_route, after_route), (wrapper, before_wrapper, after_wrapper)]:
|
||
|
|
compile(after, str(path), 'exec')
|
||
|
|
if args.check:
|
||
|
|
if before != after:
|
||
|
|
raise RuntimeError(f'Missing patch: {path}')
|
||
|
|
elif before != after:
|
||
|
|
backup = path.with_name(path.name + '.bak-image-build')
|
||
|
|
if backup.exists():
|
||
|
|
raise RuntimeError(f'Backup already exists: {backup}')
|
||
|
|
backup.write_text(before)
|
||
|
|
temp = path.with_name(path.name + '.image-build.tmp')
|
||
|
|
temp.write_text(after); temp.chmod(path.stat().st_mode & 0o777); temp.replace(path)
|
||
|
|
print('SHA256', hashlib.sha256(after.encode()).hexdigest(), path)
|
||
|
|
print('VERSIONS', json.dumps(versions))
|
||
|
|
print('ALL_PATCH_CHECKS_PASS')
|