[FIX] add sglang launch wrapper

sglang/glm53_patch.py

@@ -0,0 +1,98 @@
+#!/usr/bin/env python
+import sys
+import regex as re
+
+import sglang.srt.layers.quantization.modelopt_quant as moq
+import sglang.srt.layers.quantization.utils as quant_utils
+from sglang.srt.layers.quantization.unquant import UnquantizedLinearMethod
+
+try:
+    _use_nvfp4_dispatch = moq._use_nvfp4_dispatch
+except AttributeError:
+    try:
+        from sglang.srt.layers.quantization.modelopt_quant import _use_nvfp4_dispatch
+    except ImportError:
+        _use_nvfp4_dispatch = lambda: False
+
+GLM5_FUSED_MAP = {
+    "fused_qkv_a_proj_with_mqa": ["q_a_proj", "kv_a_proj_with_mqa"],
+    "fused_qkvbfg_a_proj": ["q_proj", "k_proj", "v_proj", "b_proj", "f_a_proj", "g_a_proj"],
+    "fused_fg_b_proj": ["f_b_proj", "g_b_proj"],
+    "qkv_proj": ["q_proj", "k_proj", "v_proj"],
+    "qkv_conv1d": ["q_conv1d", "k_conv1d", "v_conv1d"],
+    "gate_up_proj": ["gate_proj", "up_proj"],
+}
+
+if not isinstance(quant_utils._FALLBACK_FUSED_SHARDS, dict):
+    quant_utils._FALLBACK_FUSED_SHARDS = dict(quant_utils._FALLBACK_FUSED_SHARDS)
+quant_utils._FALLBACK_FUSED_SHARDS.update(GLM5_FUSED_MAP)
+
+_orig_init = moq.ModelOptQuantConfig.__init__
+def _patched_init(self, *args, **kwargs):
+    _orig_init(self, *args, **kwargs)
+    if self.packed_modules_mapping:
+        for k, v in GLM5_FUSED_MAP.items():
+            self.packed_modules_mapping.setdefault(k, v)
+    else:
+        self.packed_modules_mapping = dict(GLM5_FUSED_MAP)
+    if self.exclude_modules:
+        mapped = []
+        for name in self.exclude_modules:
+            mapped.append(name)
+            if name.startswith("model.language_model."):
+                mapped.append(name.replace("model.language_model.", "model."))
+            elif name.startswith("model.visual"):
+                mapped.append(name.replace("model.visual", "visual"))
+        self.exclude_modules = list(dict.fromkeys(mapped))
+moq.ModelOptQuantConfig.__init__ = _patched_init
+
+_orig_is_excluded = moq.ModelOptQuantConfig.is_layer_excluded
+def _patched_is_excluded(self, prefix):
+    if not self.exclude_modules:
+        return False
+    prefixes = [prefix]
+    if prefix.startswith("language_model."):
+        prefixes.append(prefix.removeprefix("language_model."))
+    head, _, tail = prefix.rpartition(".")
+    packed = self.packed_modules_mapping or {}
+    if tail in packed:
+        for shard in packed[tail]:
+            exp = f"{head}.{shard}" if head else shard
+            if exp not in prefixes:
+                prefixes.append(exp)
+    fused = {"q_a_proj", "q_b_proj", "kv_a_proj_with_mqa", "kv_b_proj"}
+    for pattern in self.exclude_modules:
+        rx = pattern.replace(".", r"\.").replace("*", r".*")
+        for pfx in prefixes:
+            if re.fullmatch(rx, pfx):
+                return True
+            for part in pfx.split("."):
+                if re.fullmatch(rx, part):
+                    return True
+        pt = pattern.rsplit(".", maxsplit=1)[-1]
+        if pt in fused:
+            for pfx in prefixes:
+                if pt in pfx.rsplit(".", maxsplit=1)[-1]:
+                    return True
+    return False
+moq.ModelOptQuantConfig.is_layer_excluded = _patched_is_excluded
+
+def _force_unquant(get_method):
+    def _patched_get(self, layer, prefix):
+        if "visual." in prefix:
+            return UnquantizedLinearMethod()
+        return get_method(self, layer, prefix)
+    return _patched_get
+
+for cls in [moq.ModelOptFp4Config, moq.ModelOptFp8Config, moq.ModelOptQuantConfig]:
+    if hasattr(cls, "get_quant_method"):
+        cls.get_quant_method = _force_unquant(cls.get_quant_method)
+
+from sglang.srt.plugins import load_plugins
+from sglang.launch_server import run_server
+from sglang.srt.server_args import prepare_server_args
+
+if __name__ == "__main__":
+    load_plugins()
+    server_args = prepare_server_args(sys.argv[1:])
+    run_server(server_args)