Where's sglang? Do you have a plan for adding Sglang support?

#2
by IlyaTers - opened

Where's sglang? Do you have a plan for adding Sglang support?

GLM-5.3-Flash-NVFP4 — sglang launch wrapper

Run nvidia/GLM-5.3-Flash-NVFP4 on sglang.

TL;DR — take the official FP8 cookbook launch command, use the NVFP4 checkpoint, run it through this wrapper. That's it.

#!/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)

Launch:

python glm53_patch.py \
  --model-path nvidia/GLM-5.3-Flash-NVFP4 \
  --tp-size 4 \
  --ep-size 4 \
  --quantization modelopt_fp4 \
  --moe-runner-backend flashinfer_cutlass \
  --trust-remote-code

All normal sglang.launch_server args are passed through after the script name.

What the wrapper fixes

  1. Shape mismatch on load ([8192, 256] vs [8192, 512]) — the checkpoint's exclude list uses model.language_model.* / model.visual* prefixes, but sglang renames them only after model construction, so excluded layers were built FP4-packed. The patch pre-maps the exclude list and forces UnquantizedLinearMethod for the vision tower.
  2. GLM-5 fused modules missing from packed_modules_mappingfused_qkvbfg_a_proj, qkv_proj, qkv_conv1d, etc. are injected.
  3. Version drift_use_nvfp4_dispatch, _input_scale_to_local_experts, is_flashinfer_megamoe don't exist in some sglang builds; wrapped in try/except / getattr.

Notes

  • moe-runner-backend flashinfer_cutlass is required
  • Verified on 4×B200, TP=4, EP=4

Sign up or log in to comment