| |
| """Apply the Agnes patch to an sglang python package directory (the one that |
| contains `srt/`). Idempotent. |
| |
| 1. srt/configs/agnes.py new file |
| 2. srt/utils/hf_transformers/common.py register AgnesConfig in _CONFIG_REGISTRY |
| 3. srt/models/qwen3_5.py load_weights: translate the Agnes checkpoint |
| (tensor prefixes, parallel-FFN fold) |
| |
| Usage: apply_patch.py <path/to/sglang> e.g. .../site-packages/sglang |
| """ |
|
|
| import os |
| import shutil |
| import sys |
|
|
| HERE = os.path.dirname(os.path.abspath(__file__)) |
| MARK = "# === agnes ===" |
|
|
| TRANSLATE = ''' |
| # === agnes === |
| # Agnes 3.0 Flash checkpoints (config model_type "agnes") use their own tensor |
| # prefixes and carry a parallel FFN branch per layer. This generator sits at |
| # the top of the weight stream and turns it into what the implementation below |
| # expects: delta_attn -> linear_attn, global_attn -> self_attn, and the branch |
| # concatenated onto the main gate / up (dim 0) and down (dim 1) projections, |
| # matching the widened intermediate_size set by sglang.srt.configs.agnes. |
| import json as _agnes_json |
| import os as _agnes_os |
| import re as _agnes_re |
| |
| _AGNES_MLP_RE = _agnes_re.compile(r"^(.*\\.layers\\.\\d+\\.mlp\\.)(gate_proj|up_proj|down_proj)\\.weight$") |
| |
| |
| class _AgnesBranchReader: |
| def __init__(self, model_path): |
| from safetensors import safe_open |
| |
| self._open = safe_open |
| self.path = model_path |
| index = _agnes_os.path.join(model_path, "model.safetensors.index.json") |
| self.weight_map = _agnes_json.load(open(index))["weight_map"] |
| self.handles = {} |
| |
| def get(self, key): |
| fn = self.weight_map[key] |
| if fn not in self.handles: |
| self.handles[fn] = self._open(_agnes_os.path.join(self.path, fn), framework="pt", device="cpu") |
| return self.handles[fn].get_tensor(key) |
| |
| |
| def _agnes_translate_weights(model, weights): |
| cfg = getattr(model.config, "text_config", None) or model.config |
| width = int(getattr(cfg, "agnes_parallel_ffn_intermediate_size", 0) or 0) |
| if width <= 0: |
| yield from weights |
| return |
| model_path = ( |
| getattr(cfg, "agnes_model_path", None) |
| or getattr(model.config, "agnes_model_path", None) |
| or _agnes_os.environ.get("AGNES_MODEL_PATH") |
| or getattr(model.config, "_name_or_path", None) |
| ) |
| if not model_path or not _agnes_os.path.isdir(model_path): |
| raise RuntimeError( |
| f"agnes: cannot locate the checkpoint directory (got {model_path!r}); " |
| "set AGNES_MODEL_PATH to the model directory" |
| ) |
| reader = _AgnesBranchReader(model_path) |
| for name, w in weights: |
| if ".mlp.parallel_ffn." in name: |
| continue |
| m = _AGNES_MLP_RE.match(name) |
| if m and "visual" not in name and not name.startswith("mtp"): |
| extra = reader.get(f"{m.group(1)}parallel_ffn.{m.group(2)}.weight") |
| dim = 1 if m.group(2) == "down_proj" else 0 |
| w = torch.cat([w, extra.to(device=w.device, dtype=w.dtype)], dim=dim) |
| name = name.replace(".delta_attn.", ".linear_attn.").replace(".global_attn.", ".self_attn.") |
| yield name, w |
| # === /agnes === |
| |
| ''' |
|
|
| REGISTER = ''' |
| # === agnes === |
| from sglang.srt.configs.agnes import AgnesConfig as _AgnesConfig |
| |
| _CONFIG_REGISTRY[_AgnesConfig.model_type] = _AgnesConfig |
| ''' |
|
|
|
|
| def patch_file(path, edit): |
| src = open(path, encoding="utf-8").read() |
| if MARK in src: |
| return "already patched" |
| out = edit(src) |
| if out is None: |
| raise SystemExit(f"anchor not found in {path}") |
| open(path, "w", encoding="utf-8").write(out) |
| return "patched" |
|
|
|
|
| def edit_model_file(src): |
| anchor = "QWEN3_5_KV_SCALE_MAPPER = WeightsMapper(" |
| hook = " weights = QWEN3_5_KV_SCALE_MAPPER.apply(weights)\n" |
| if anchor not in src or src.count(hook) < 1: |
| return None |
| src = src.replace(anchor, TRANSLATE + anchor, 1) |
| src = src.replace(hook, " weights = _agnes_translate_weights(self, weights)\n" + hook) |
| return src |
|
|
|
|
| def main(): |
| if len(sys.argv) != 2: |
| sys.exit(__doc__) |
| pkg = os.path.abspath(sys.argv[1]) |
| srt = os.path.join(pkg, "srt") |
| if not os.path.isdir(srt): |
| sys.exit(f"{pkg} does not contain srt/") |
| dst = os.path.join(srt, "configs", "agnes.py") |
| shutil.copy2(os.path.join(HERE, "agnes_sglang_config.py"), dst) |
| print(f"configs/agnes.py: installed") |
| print("utils/hf_transformers/common.py:", patch_file( |
| os.path.join(srt, "utils", "hf_transformers", "common.py"), lambda s: s.rstrip("\n") + "\n" + REGISTER)) |
| print("models/qwen3_5.py:", patch_file(os.path.join(srt, "models", "qwen3_5.py"), edit_model_file)) |
| print("APPLY_PATCH_OK") |
|
|
|
|
| if __name__ == "__main__": |
| main() |
|
|
