first commit

routing_fix/patch_sigmoid_wrapper.py

@@ -0,0 +1,62 @@
+"""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')