#!/usr/bin/env python3
"""Teach the glm5next model classes their fused projection names.

local-inference-lab/GLM-5.3-Flash-NVFP4-4p67 is the first checkpoint of this
family that QUANTIZES the individual attention projections (per-shard
MXFP8 entries `self_attn.q_proj` … in hf_quant_config). vLLM fuses those
shards into single modules (`in_proj_qkvbfg_a`, `fused_qkv_a_proj`,
`wk_weights_proj`), and the ModelOpt MIXED_PRECISION config resolves a
fused module's quant_algo through its shards ONLY via the model class's
`packed_modules_mapping` (handed to the quant config by
model_executor/model_loader/utils.py). Without it every fused linear comes
up UnquantizedLinearMethod and weight loading dies with

    KeyError: 'layers.0.self_attn.in_proj_qkvbfg_a.weight_scale'

Anchors self-assert; idempotent.
"""

import py_compile
import sys

B = "/usr/local/lib/python3.12/dist-packages/vllm/models/glm5next/nvidia"

MAPPING = '''
    # NVFP4-4p67: quantized per-shard attention/shared-expert projections
    # are fused by vLLM; expose the shard names so the ModelOpt
    # MIXED_PRECISION config can resolve the fused module's quant_algo
    # (see model_loader/utils.py packed_modules_mapping pass-through).
    packed_modules_mapping = {
        "gate_up_proj": ["gate_proj", "up_proj"],
        "in_proj_qkvbfg_a": ["q_proj", "k_proj", "v_proj",
                             "b_proj", "f_a_proj", "g_a_proj"],
        "fused_qkv_a_proj": ["q_a_proj", "kv_a_proj_with_mqa"],
        "wk_weights_proj": ["wk", "weights_proj"],
    }
'''

TARGETS = [
    (
        f"{B}/model.py",
        "class Glm5NextForCausalLM(\n"
        "    nn.Module, HasInnerState, SupportsPP, MixtureOfExperts, IsHybrid\n"
        "):\n",
    ),
    (
        f"{B}/mtp.py",
        "class Glm5NextMTP(nn.Module, DeepseekV2MixtureOfExperts):\n",
    ),
]

for path, anchor in TARGETS:
    src = open(path).read()
    if "packed_modules_mapping" in src:
        print(f"already patched: {path}")
        continue
    n = src.count(anchor)
    assert n == 1, f"anchor {n}x in {path} — day-0 image drifted, review me"
    open(path, "w").write(src.replace(anchor, anchor + MAPPING, 1))
    py_compile.compile(path, doraise=True)
    print(f"patched: {path}")
print("patch_4p67 OK")