first commit

image_build/apply_patches.py

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