#!/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"],
    }
'''

# CRITICAL: it must live on Glm5NextForConditionalGeneration (the registered
# class). configure_quant_config reads BOTH hf_to_vllm_mapper and
# packed_modules_mapping off that class; Glm4v's inherited mapper rewrites
# the config's quantized_layers keys into vLLM spelling, and a class-level
# mapping on the inner CausalLM would SHADOW the inherited Glm4v mapper's
# companion config path. MTP is constructed as a separate root model, so it
# needs its own copy.
MAPPING_TARGETS = [
    (
        f"{B}/model.py",
        "class Glm5NextForCausalLM(\n"
        "    nn.Module, HasInnerState, SupportsPP, MixtureOfExperts, IsHybrid\n"
        "):\n",
    ),
    (
        f"{B}/model.py",
        "class Glm5NextForConditionalGeneration(\n"
        "    Glm4vForConditionalGeneration, HasInnerState, IsHybrid\n"
        "):\n",
    ),
    (
        f"{B}/mtp.py",
        "class Glm5NextMTP(nn.Module, DeepseekV2MixtureOfExperts):\n",
    ),
]

for path, anchor in MAPPING_TARGETS:
    src = open(path).read()
    # count INSERTED BLOCKS (marker appears once per patch), not the plain
    # string (which also appears in comments and inherited bases).
    want = sum(1 for p, _ in MAPPING_TARGETS if p == path)
    if src.count('"wk_weights_proj": ["wk", "weights_proj"],') >= want:
        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}")
# KDA strip gate. day-0 KDA construction FORCIBLY unquantizes its
# projections ("BF16 even in fp8 checkpoints") by nulling
# vllm_config.quant_config around super().__init__ — which is exactly the
# 4p67 assumption-break: this checkpoint stores MXFP8 scales for
# self_attn.{q,k,v,b,f_a,g_a}_proj, so the fused linear comes up BF16 and
# the loader dies routing those scales. Make the strip conditional on the
# checkpoint genuinely lacking KDA shard scales.
KDA_STRIP_OLD = '''        saved_quant_config = vllm_config.quant_config
        vllm_config.quant_config = None
        super().__init__(config, vllm_config, prefix)
        vllm_config.quant_config = saved_quant_config
'''

KDA_STRIP_NEW = '''        saved_quant_config = vllm_config.quant_config
        _qc = vllm_config.quant_config
        _kda_quantized = _qc is not None and any(
            ("%s.%s" % (prefix, _s)) in getattr(_qc, "quantized_layers", {})
            for _s in (
                "q_proj", "k_proj", "v_proj",
                "b_proj", "f_a_proj", "g_a_proj",
            )
        )
        if not _kda_quantized:
            # day-0 checkpoints keep the KDA projections BF16 (no scales
            # stored for them); NVFP4-4p67 stores MXFP8 -> keep quant.
            vllm_config.quant_config = None
        super().__init__(config, vllm_config, prefix)
        vllm_config.quant_config = saved_quant_config
'''

kda_path = f"{B}/kda.py"
src = open(kda_path).read()
if "_kda_quantized" in src:
    print(f"already patched: {kda_path}")
else:
    n = src.count(KDA_STRIP_OLD)
    assert n == 1, f"KDA strip anchor {n}x in {kda_path} — review me"
    open(kda_path, "w").write(src.replace(KDA_STRIP_OLD, KDA_STRIP_NEW, 1))
    py_compile.compile(kda_path, doraise=True)
    print(f"patched: {kda_path}")

# (2) Indexer wk_weights_proj: attention.py hardcodes quant_config=None
# (legacy path dequantizes FP8 wk to BF16 on load). Gate it the same way as
# KDA: if the checkpoint declares MXFP8 for indexer.wk / indexer.weights_proj
# (NVFP4-4p67), build the fused linear QUANTIZED and let the stacked loader
# route weights+scales.
IDX_OLD = '''        self.wk_weights_proj = MergedColumnParallelLinear(
            hidden_size,
            [self.head_dim, self.n_head],
            bias=False,
            quant_config=None,
'''
IDX_NEW = '''        _idx_quantized = quant_config is not None and any(
            ("%s.%s" % (prefix, _s)) in getattr(quant_config, "quantized_layers", {})
            for _s in ("wk", "weights_proj")
        )
        self.wk_weights_proj = MergedColumnParallelLinear(
            hidden_size,
            [self.head_dim, self.n_head],
            bias=False,
            quant_config=(quant_config if _idx_quantized else None),
'''

apath = f"{B}/attention.py"
src = open(apath).read()
if "_idx_quantized" in src:
    print(f"already patched: {apath}")
else:
    n = src.count(IDX_OLD)
    assert n == 1, f"indexer anchor {n}x in {apath} — review me"
    open(apath, "w").write(src.replace(IDX_OLD, IDX_NEW, 1))
    py_compile.compile(apath, doraise=True)
    print(f"patched: {apath}")

# (3) Loader shims: legacy FP8 dequant helpers must not eat MXFP8 tensors
# (4p67 names scales `weight_scale` — the wk shim's is_weight test would
# swallow the fp8 weight and then drop the scale). Skip the shims when the
# target fused param is quantized (has a weight_scale param).
model_path = f"{B}/model.py"
src = open(model_path).read()
WK_OLD = '''def _try_load_fp8_indexer_wk(name, tensor, buf, params_dict, loaded_params):
    if "indexer.wk." not in name or "wk_weights" in name:
        return False
'''
WK_NEW = '''def _try_load_fp8_indexer_wk(name, tensor, buf, params_dict, loaded_params):
    if "indexer.wk." not in name or "wk_weights" in name:
        return False
    if (f"{name.rsplit('.wk.', 1)[0]}.wk_weights_proj.weight_scale"
            in params_dict):
        # MXFP8 checkpoint: fused param is quantized, stacked loader owns it.
        return False
'''
ATTN_OLD = '''    target_w = f"{layer_prefix}.{target_base}.weight"
    target_s = f"{layer_prefix}.{target_base}.weight_scale_inv"
'''
ATTN_NEW = '''    target_w = f"{layer_prefix}.{target_base}.weight"
    target_s = f"{layer_prefix}.{target_base}.weight_scale_inv"
    if f"{layer_prefix}.{target_base}.weight_scale" in params_dict:
        # MXFP8 checkpoint: param is quantized, stacked loader owns it.
        return False
'''
if "stacked loader owns it" in src:
    print(f"already patched: {model_path} (wk shim)")
else:
    n = src.count(WK_OLD)
    assert n == 1, f"wk-shim anchor {n}x in {model_path} — review me"
    src = src.replace(WK_OLD, WK_NEW, 1)
    open(model_path, "w").write(src)
    py_compile.compile(model_path, doraise=True)
    print(f"patched: {model_path} (wk shim)")

if "MXFP8 checkpoint: param is quantized" in src:
    print(f"already patched: {model_path} (attn shim)")
else:
    src = open(model_path).read()
    n = src.count(ATTN_OLD)
    assert n == 1, f"attn-shim anchor {n}x in {model_path} — review me"
    src = src.replace(ATTN_OLD, ATTN_NEW, 1)
    open(model_path, "w").write(src)
    py_compile.compile(model_path, doraise=True)
    print(f"patched: {model_path} (attn shim)")

def helper(name, uses):
    h = (
        "def _glm5_q_layers(quant_config):\n"
        "    # quantized_layers keys are HF-spelled at module-construction\n"
        "    # time and vLLM-spelled after apply_vllm_mapper; test both.\n"
        "    return set(getattr(quant_config, 'quantized_layers', {}) or {})\n"
        "\n"
        "\n"
        "def _glm5_quantized(quant_config, prefix, shards):\n"
        "    layers = _glm5_q_layers(quant_config)\n"
        "    if not layers:\n"
        "        return False\n"
        "    tail = prefix\n"
        "    for lead in (\n"
        "        'language_model.model.', 'model.language_model.',\n"
        "        'model.', 'language_model.',\n"
        "    ):\n"
        "        if tail.startswith(lead):\n"
        "            tail = tail[len(lead):]\n"
        "            break\n"
        "    return any(\n"
        "        k.endswith(f'.{tail}.{s}') or k == f'{tail}.{s}'\n"
        "        for s in shards\n"
        "        for k in layers\n"
        "    )\n\n\n"
    )
    # place right before the first class that needs it
    anchor = uses
    return h, anchor


for path, defs_marker, insert_before in [
    (f"{B}/kda.py", "_glm5_quantized", "class Glm5NextLinearAttention"),
    (f"{B}/attention.py", "_glm5_quantized", "class Indexer(nn.Module):"),
]:
    src = open(path).read()
    if defs_marker in src:
        print(f"helper already in {path}")
        continue
    idx = src.index(insert_before)
    helper_src, _ = helper(path, insert_before)
    src = src[:idx] + helper_src + src[idx:]
    open(path, "w").write(src)
    py_compile.compile(path, doraise=True)
    print(f"helper added: {path}")

# rewrite the two constructor gates to use the both-spellings helper
GATES = [
    (
        f"{B}/kda.py",
        '''        _kda_quantized = _qc is not None and any(
            ("%s.%s" % (prefix, _s)) in getattr(_qc, "quantized_layers", {})
            for _s in (
                "q_proj", "k_proj", "v_proj",
                "b_proj", "f_a_proj", "g_a_proj",
            )
        )
''',
        "        _kda_quantized = _glm5_quantized(\n"
        "            _qc, prefix,\n"
        '            ("q_proj", "k_proj", "v_proj",\n'
        '             "b_proj", "f_a_proj", "g_a_proj"),\n'
        "        )\n",
    ),
    (
        f"{B}/attention.py",
        '''        _idx_quantized = quant_config is not None and any(
            ("%s.%s" % (prefix, _s)) in getattr(quant_config, "quantized_layers", {})
            for _s in ("wk", "weights_proj")
        )
''',
        "        _idx_quantized = _glm5_quantized(\n"
        '            quant_config, prefix, ("wk", "weights_proj")\n'
        "        )\n"
        "        import logging as _lg\n"
        '        _lg.getLogger("vllm").warning(\n'
        '            "4p67 IDX GATE prefix=%s q=%s qc=%s n=%s",\n'
        '            prefix, _idx_quantized,\n'
        "            type(quant_config).__name__,\n"
        '            len(getattr(quant_config, "quantized_layers", {}) or {}),\n'
        "        )\n",
    ),
]
for path, old, new in GATES:
    src = open(path).read()
    if "_glm5_quantized(" in src.split("def _glm5_quantized", 1)[-1]:
        print(f"gate already rewritten: {path}")
        continue
    n = src.count(old)
    assert n == 1, f"gate anchor {n}x in {path} — review me"
    open(path, "w").write(src.replace(old, new, 1))
    py_compile.compile(path, doraise=True)
    print(f"gate rewritten: {path}")

# (5) DecoderLayer hardcodes quant_config=None for MLA attention ("BF16 in
# checkpoint" — again broken by 4p67). Gate it; helper lives in kda.py
# (always importable; model.py already imports from it? if not, guard).
MLA_OLD = "                quant_config=None,  # MLA projections are BF16 in checkpoint\n"
MLA_NEW = """                quant_config=(
                    quant_config
                    if _glm5_quantized(
                        quant_config,
                        f"{prefix}.self_attn",
                        ("q_a_proj", "kv_a_proj_with_mqa"),
                    )
                    else None
                ),  # 4p67: quantized when checkpoint stores MLA scales
"""
mla_path = f"{B}/model.py"
src = open(mla_path).read()
if "4p67: quantized when checkpoint stores MLA scales" in src:
    print(f"already patched: {mla_path} (mla gate)")
else:
    n = src.count(MLA_OLD)
    assert n == 1, f"mla-gate anchor {n}x in {mla_path} — review me"
    src = src.replace(MLA_OLD, MLA_NEW, 1)
    if "from .kda import" in src:
        import re as _re

        src = _re.sub(
            r"from \.kda import (\w+)",
            r"from .kda import \1, _glm5_quantized",
            src,
            count=1,
        )
    else:
        raise AssertionError(
            f"no 'from .kda import' in {mla_path} — add import explicitly"
        )
    assert "_glm5_quantized" in "\n".join(src.splitlines()[:120]), "import missing"
    open(mla_path, "w").write(src)
    py_compile.compile(mla_path, doraise=True)
    print(f"patched: {mla_path} (mla gate)")

print("patch_4p67 OK")