63 lines
2.7 KiB
Python
63 lines
2.7 KiB
Python
|
"""Run inside srx_test: fix lower-level keyword without changing Torch schema."""
|
||
|
|
import ast
|
||
|
|
import copy
|
||
|
|
import datetime
|
||
|
|
import inspect
|
||
|
|
from pathlib import Path
|
||
|
|
from types import SimpleNamespace
|
||
|
|
|
||
|
|
import kunlun_ops
|
||
|
|
|
||
|
|
path = Path('/opt/vllm_kunlun/lib/python3.10/site-packages/vllm_kunlun/ops/_custom_ops.py')
|
||
|
|
original = path.read_text()
|
||
|
|
signature = inspect.signature(kunlun_ops.moe_sigmoid_group_topk_norm)
|
||
|
|
assert 'block_statistic' in signature.parameters
|
||
|
|
names = {'moe_sigmoid_group_topk_norm', 'moe_sigmoid_group_topk_norm_cuda'}
|
||
|
|
tree = ast.parse(original)
|
||
|
|
functions = [n for n in tree.body if isinstance(n, ast.FunctionDef) and n.name in names]
|
||
|
|
assert len(functions) == 2
|
||
|
|
lines = original.splitlines(keepends=True)
|
||
|
|
changed = 0
|
||
|
|
for fn in functions:
|
||
|
|
for call in ast.walk(fn):
|
||
|
|
if not isinstance(call, ast.Call) or ast.unparse(call.func) != 'kunlun_ops.moe_sigmoid_group_topk_norm':
|
||
|
|
continue
|
||
|
|
for keyword in call.keywords:
|
||
|
|
if keyword.arg == 'block_static':
|
||
|
|
index = keyword.lineno - 1
|
||
|
|
assert 'block_static=block_static' in lines[index]
|
||
|
|
lines[index] = lines[index].replace('block_static=block_static', 'block_statistic=block_static')
|
||
|
|
changed += 1
|
||
|
|
assert changed in (0, 2), f'Unexpected matching calls: {changed}'
|
||
|
|
patched = ''.join(lines)
|
||
|
|
compile(patched, str(path), 'exec')
|
||
|
|
# Execute both actual wrapper bodies against the installed lower-level signature.
|
||
|
|
for fn in [n for n in ast.parse(patched).body if isinstance(n, ast.FunctionDef) and n.name in names]:
|
||
|
|
node = copy.deepcopy(fn)
|
||
|
|
node.decorator_list = []
|
||
|
|
node.returns = None
|
||
|
|
for arg in node.args.args:
|
||
|
|
arg.annotation = None
|
||
|
|
calls = []
|
||
|
|
def spy(*args, **kwargs):
|
||
|
|
bound = signature.bind(*args, **kwargs)
|
||
|
|
calls.append(bound.arguments)
|
||
|
|
scope = {'kunlun_ops': SimpleNamespace(moe_sigmoid_group_topk_norm=spy)}
|
||
|
|
exec(compile(ast.fix_missing_locations(ast.Module(body=[node], type_ignores=[])), '<wrapper-test>', 'exec'), scope)
|
||
|
|
sentinels = {name: object() for name in ('x', 'topk_index', 'norm_score', 'block_static', 'bias')}
|
||
|
|
scope[fn.name](**sentinels, scale=1.0, n_group=8, topk_group=4)
|
||
|
|
assert calls[0]['block_statistic'] is sentinels['block_static']
|
||
|
|
assert calls[0]['bias'] is sentinels['bias']
|
||
|
|
print('WRAPPER_SIGNATURE_PASS', fn.name)
|
||
|
|
if changed:
|
||
|
|
backup = path.with_name(path.name + '.bak-sigmoid-keyword-' + datetime.datetime.now().strftime('%Y%m%dT%H%M%S%f'))
|
||
|
|
backup.write_text(original)
|
||
|
|
temp = path.with_name(path.name + '.routing-fix.tmp')
|
||
|
|
temp.write_text(patched)
|
||
|
|
temp.chmod(path.stat().st_mode & 0o777)
|
||
|
|
temp.replace(path)
|
||
|
|
print('BACKUP', backup)
|
||
|
|
print('PATCHED', path)
|
||
|
|
else:
|
||
|
|
print('ALREADY_PATCHED')
|