first commit

routing_fix/patch_moe_routing.py

@@ -0,0 +1,141 @@
+#!/usr/bin/env python3
+"""Patch the paused/running container's installed Kunlun INT8 MoE, with backup.
+
+Run on the Docker host. Tests exercise the actual method through routing dispatch
+with tensor allocation/kernel calls mocked; they do not execute XPU kernels.
+"""
+import ast
+import copy
+import datetime
+import hashlib
+import io
+import subprocess
+import tarfile
+from types import SimpleNamespace as NS
+
+CONTAINER = 'srx_test'
+DIRECTORY = '/opt/vllm_kunlun/lib/python3.10/site-packages/vllm_kunlun/quantization/compressed_tensors'
+NAME = 'compressed_tensors_moe.py'
+CLASS = 'KunlunCompressedTensorsW8A8Int8MoEMethod'
+MARKER = '# Kunlun vLLM 0.25 routing compatibility fix'
+INSERT = '''        # Kunlun vLLM 0.25 routing compatibility fix
+        # RoutedExperts.forward_monolithic passes only layer/x/logits/input_ids.
+        # Read routing from the layer, not legacy method defaults. In particular,
+        # routed_scaling_factor may be 1 because the runner applies it externally.
+        scoring_func = layer.scoring_func
+        num_expert_group = layer.num_expert_group
+        topk_group = layer.topk_group
+        e_score_correction_bias = layer.e_score_correction_bias
+        routed_scaling_factor = layer.routed_scaling_factor
+        if scoring_func not in ("softmax", "sigmoid"):
+            raise NotImplementedError(f"Unsupported Kunlun INT8 MoE routing: {scoring_func}")
+
+'''
+
+
+def method_node(source):
+    module = ast.parse(source)
+    cls = next(n for n in module.body if isinstance(n, ast.ClassDef) and n.name == CLASS)
+    return next(n for n in cls.body if isinstance(n, ast.FunctionDef) and n.name == 'apply_monolithic')
+
+
+def patch_source(source):
+    if MARKER in source:
+        return source
+    node = method_node(source)
+    first = node.body[0]
+    if ast.unparse(first) != 'hidden_states = x':
+        raise RuntimeError('Unexpected installed method; refusing to patch')
+    lines = source.splitlines(keepends=True)
+    lines.insert(first.lineno - 1, INSERT)
+    patched = ''.join(lines)
+    compile(patched, NAME, 'exec')
+    return patched
+
+
+class Captured(Exception):
+    pass
+
+
+def dispatch(source, scoring, groups, top_groups, scale, bias):
+    node = copy.deepcopy(method_node(source))
+    for arg in node.args.posonlyargs + node.args.args + node.args.kwonlyargs:
+        arg.annotation = None
+    node.returns = None
+    calls = []
+
+    def capture(kind):
+        def callback(**kwargs):
+            calls.append((kind, kwargs))
+            raise Captured()
+        return callback
+
+    ops = NS(moe_softmax_topk_norm=capture('softmax'),
+             moe_sigmoid_group_topk_norm=capture('sigmoid'))
+    torch = NS(float32='float32', int32='int32',
+               empty=lambda *a, **k: object(), zeros=lambda *a, **k: object(), ops=NS(_C=ops))
+    scope = {'torch': torch}
+    exec(compile(ast.fix_missing_locations(ast.Module(body=[node], type_ignores=[])), '<actual-method>', 'exec'), scope)
+    layer = NS(scoring_func=scoring, num_expert_group=groups, topk_group=top_groups,
+               routed_scaling_factor=scale, e_score_correction_bias=bias,
+               w13_weight=NS(shape=(256, 512, 128)), w2_weight=NS(shape=(256, 128, 256)))
+    try:
+        scope['apply_monolithic'](NS(moe=NS(experts_per_token=8)), layer,
+                                  NS(shape=(1, 128), device='cpu'), NS(float=lambda: object()),
+                                  input_ids=None)
+    except Captured:
+        pass
+    assert len(calls) == 1
+    return calls[0]
+
+
+def regression_tests(original, patched):
+    bias = object()
+    if MARKER not in original:
+        kind, _ = dispatch(original, 'sigmoid', 8, 4, 1.0, bias)
+        assert kind == 'softmax', 'Expected to reproduce the old routing bug'
+        print('REPRODUCED: old method used softmax for DeepSeek sigmoid routing')
+    for scale in (1.0, 2.5):
+        kind, kwargs = dispatch(patched, 'sigmoid', 8, 4, scale, bias)
+        assert kind == 'sigmoid'
+        assert kwargs['n_group'] == 8 and kwargs['topk_group'] == 4
+        assert kwargs['bias'] is bias and kwargs['scale'] == scale
+    assert dispatch(patched, 'softmax', None, None, 1.0, None)[0] == 'softmax'
+    kind, kwargs = dispatch(patched, 'sigmoid', None, None, 1.0, bias)
+    assert kind == 'sigmoid' and kwargs['bias'] is bias
+    print('PASS: sigmoid groups/bias, external/internal scaling, softmax and ungrouped sigmoid')
+
+
+def read_installed():
+    data = subprocess.check_output(['docker', 'cp', f'{CONTAINER}:{DIRECTORY}/{NAME}', '-'])
+    with tarfile.open(fileobj=io.BytesIO(data)) as archive:
+        entry = next(m for m in archive.getmembers() if m.isfile())
+        return entry, archive.extractfile(entry).read()
+
+
+def main():
+    entry, original_bytes = read_installed()
+    original = original_bytes.decode()
+    patched = patch_source(original)
+    regression_tests(original, patched)
+    if patched == original:
+        print('ALREADY_PATCHED')
+        return
+    backup = NAME + '.bak-routing-' + datetime.datetime.now().strftime('%Y%m%dT%H%M%S%f')
+    output = io.BytesIO()
+    with tarfile.open(fileobj=output, mode='w') as archive:
+        for name, content in [(backup, original_bytes), (NAME, patched.encode())]:
+            info = copy.copy(entry)
+            info.name, info.size = name, len(content)
+            archive.addfile(info, io.BytesIO(content))
+    subprocess.run(['docker', 'cp', '-', f'{CONTAINER}:{DIRECTORY}'], input=output.getvalue(), check=True)
+    _, installed = read_installed()
+    assert installed == patched.encode(), 'Installed content differs from patch'
+    print('INSTALLED', DIRECTORY + '/' + NAME)
+    print('BACKUP', DIRECTORY + '/' + backup)
+    print('SHA256', hashlib.sha256(installed).hexdigest())
+    print('Restart the vLLM service to load the fix; existing workers retain old code.')
+
+
+if __name__ == '__main__':
+    main()

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')