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