code/tools/collect_expert_stats_v2.py

#!/usr/bin/env python3
"""Per-expert routing stats on the STOCK NVFP4 GLM-5.2, coding/agentic corpus.

Runs /tmp/glm52-dl/nvfp4-full (all experts NVFP4, near-lossless routing)
across 8 GPUs with a forced modular MoE backend so RoutedExperts sees
topk_ids (VLLM_EXPERT_STATS_DIR hook in routed_experts.py).

Corpus (200 prompts, up to ~7.9k tokens each):
  30% raw code continuation       (vLLM py / cu / cpp sources)
  12% bash / automation corpus    (real CI scripts, bash completions,
                                   Dockerfiles, init scripts)
  34% multi-turn tool-calling     (GLM chat template + tool schemas; tool
                                   results are real logs, diffs, file
                                   contents, tracebacks; coding + devops +
                                   desktop-automation tasks)
  10% coding chat                 (debug/review/explain, embedded code)
  11% medical                     (MedQA textbook continuation + clinical
                                   knowledge Q&A -- generic knowledge only)
  3%  general chat

Output: /data/glm52-expert-stats-v2.npz  (counts [n_layers, 256])
"""
import glob
import json
import os
import random
import subprocess

import numpy as np

MODEL = "/tmp/glm52-dl/nvfp4-full"
STATS_DIR = "/data/glm52-expert-stats-v2"
OUT = "/data/glm52-expert-stats-v2.npz"
N_EXPERTS = 256
MAX_TOKENS = 7900
ROOT = "/home/coder/git/glm52"

TOOLS = [
    {"type": "function", "function": {
        "name": "bash",
        "description": "Run a shell command in the repo and return stdout/stderr.",
        "parameters": {"type": "object", "properties": {
            "command": {"type": "string", "description": "Command to run"},
            "timeout": {"type": "integer", "description": "Seconds"}},
            "required": ["command"]}}},
    {"type": "function", "function": {
        "name": "read_file",
        "description": "Read a file, optionally a line range.",
        "parameters": {"type": "object", "properties": {
            "path": {"type": "string"},
            "start_line": {"type": "integer"},
            "end_line": {"type": "integer"}}, "required": ["path"]}}},
    {"type": "function", "function": {
        "name": "edit_file",
        "description": "Replace old_string with new_string in a file.",
        "parameters": {"type": "object", "properties": {
            "path": {"type": "string"},
            "old_string": {"type": "string"},
            "new_string": {"type": "string"}},
            "required": ["path", "old_string", "new_string"]}}},
    {"type": "function", "function": {
        "name": "python",
        "description": "Execute python code and return the output.",
        "parameters": {"type": "object", "properties": {
            "code": {"type": "string"}}, "required": ["code"]}}},
    {"type": "function", "function": {
        "name": "search_code",
        "description": "Regex search across the repository.",
        "parameters": {"type": "object", "properties": {
            "pattern": {"type": "string"},
            "glob": {"type": "string"}}, "required": ["pattern"]}}},
]


def _read(path, limit=20000):
    try:
        return open(path, errors="ignore").read()[:limit]
    except OSError:
        return ""


def _sh(cmd):
    try:
        return subprocess.run(
            cmd, shell=True, capture_output=True, text=True, timeout=20,
            cwd=ROOT).stdout[:20000]
    except Exception:
        return ""


def gather_artifacts(rng):
    """Real terminal/log/diff/file content to use as tool results."""
    art = {"logs": [], "code": [], "diffs": [], "listings": [], "json": []}
    for lf in glob.glob("/tmp/glm52-*.log") + glob.glob("/data/*.log"):
        txt = _read(lf, 400000)
        for i in range(0, len(txt) - 4000, 40000):
            art["logs"].append(txt[i:i + rng.randint(1500, 5000)])
    code_files = (glob.glob(f"{ROOT}/vllm/vllm/**/*.py", recursive=True)
                  + glob.glob(f"{ROOT}/vllm/csrc/**/*.cu", recursive=True)
                  + glob.glob(f"{ROOT}/vllm/csrc/**/*.cuh", recursive=True)
                  + glob.glob(f"{ROOT}/tools/*.py"))
    rng.shuffle(code_files)
    art["code_files"] = code_files
    for p in code_files[:400]:
        t = _read(p)
        if len(t) > 500:
            art["code"].append((p, t))
    d = _sh("git -C vllm diff --stat; git -C vllm diff | head -800")
    if d:
        for i in range(0, len(d), 4000):
            art["diffs"].append(d[i:i + 4000])
    art["listings"].append(_sh("ls -la vllm/vllm/model_executor/layers/quantization/ | head -40"))
    art["listings"].append(_sh("find vllm/csrc/quantization -type f | head -30"))
    art["json"].append(_read("/data/glm52-v3/config.json", 4000))
    art["json"].append(_read(f"{ROOT}/hybrid_plan.json", 4000))
    return art


USER_TASKS = [
    "The test suite is failing after my last change. Find the failing tests, "
    "diagnose the root cause, and fix it.",
    "Profile why model loading is slow and optimize the hot path.",
    "Add per-layer logging of expert routing entropy to the MoE layer, with "
    "an env var to enable it.",
    "There's a CUDA illegal memory access in the new kernel under batch>64. "
    "Find and fix it.",
    "Refactor the quantization config to support per-layer overrides without "
    "breaking existing checkpoints.",
    "Review the recent diff for correctness and potential race conditions.",
    "The server OOMs at long context. Figure out the memory budget and fix "
    "the allocation.",
    "Write a benchmark comparing the two MoE dequant paths and report results.",
    "Investigate why throughput dropped 20% between these two log files.",
    "Implement graceful fallback when the JIT kernel build fails.",
    # devops / desktop-automation / scripting tasks
    "Write a bash script that watches a downloads directory and sorts new "
    "files into folders by type, then install it as a systemd user service.",
    "Set up a cron job that backs up my projects directory to a remote host "
    "nightly with rsync, keeping 7 rotated snapshots and logging failures.",
    "My disk is nearly full. Find what's consuming space and clean it up "
    "safely, keeping anything modified in the last week.",
    "Automate my morning setup: a script that starts tmux with three named "
    "windows, activates the venv, tails the server log, and opens the todo "
    "file in the editor.",
    "Batch-rename all photos in this directory tree from IMG_xxxx to "
    "date-based names using their EXIF timestamps, handling collisions.",
    "Write a script to monitor GPU memory and temperature and alert me via "
    "desktop notification when either crosses a threshold.",
    "Migrate all my dotfiles into a git-managed directory with symlinks and "
    "an install script that works on a fresh machine.",
    "Bulk-convert these CSV exports to parquet, validate row counts match, "
    "and produce a summary report of schema differences.",
]

THINK_SNIPPETS = [
    "Let me start by looking at the failing area.",
    "I need to inspect the code first.",
    "Let me check the logs for the actual error.",
    "First I'll search for where this is defined.",
    "Let me look at the diff to understand what changed.",
    "Now let me verify the fix compiles and passes tests.",
    "The error suggests a shape mismatch; checking the tensor layout.",
    "That output confirms the hypothesis. Applying a fix.",
    "Tests pass now. Let me summarize the change.",
]


def make_agentic_session(rng, art, tok):
    msgs = [{"role": "system", "content":
             "You are a senior software engineer agent working in a large "
             "ML inference codebase. Use the available tools to investigate "
             "and make changes. Be precise and verify your work."}]
    msgs.append({"role": "user", "content": rng.choice(USER_TASKS)})
    n_calls = rng.randint(3, 7)
    cid = 0
    for _ in range(n_calls):
        cid += 1
        kind = rng.random()
        if kind < 0.35:
            cmd = rng.choice([
                "python -m pytest tests/kernels -x -q 2>&1 | tail -30",
                "grep -rn 'def apply' vllm/model_executor/layers/fused_moe/ | head",
                "git diff --stat", "nvidia-smi | head -20",
                "python tools/budget.py",
                "ninja -C build 2>&1 | tail -20",
            ])
            call = {"name": "bash", "arguments": {"command": cmd}}
            result = rng.choice(art["logs"] + art["diffs"] + art["listings"])
        elif kind < 0.6:
            p, t = rng.choice(art["code"])
            s = rng.randrange(0, max(1, len(t) - 4000))
            call = {"name": "read_file", "arguments": {
                "path": p.replace(ROOT + "/", ""),
                "start_line": s // 60, "end_line": s // 60 + 80}}
            result = t[s:s + rng.randint(1500, 4000)]
        elif kind < 0.75:
            call = {"name": "search_code", "arguments": {
                "pattern": rng.choice([
                    r"topk_ids", r"create_weights", r"kv_cache_dtype",
                    r"cudaMemcpy", r"register_parameter"]),
                "glob": "**/*.py"}}
            result = rng.choice(art["listings"] + art["diffs"])
        elif kind < 0.9:
            p, t = rng.choice(art["code"])
            s = rng.randrange(0, max(1, len(t) - 2000))
            old = t[s:s + rng.randint(80, 300)]
            call = {"name": "edit_file", "arguments": {
                "path": p.replace(ROOT + "/", ""), "old_string": old,
                "new_string": old + "  # reviewed"}}
            result = "OK: 1 replacement made."
        else:
            call = {"name": "python", "arguments": {"code":
                    "import json\ncfg=json.load(open('config.json'))\n"
                    "print(cfg['quantization_config']['quant_method'])"}}
            result = rng.choice(art["json"])
        msgs.append({"role": "assistant",
                     "content": rng.choice(THINK_SNIPPETS),
                     "tool_calls": [{"id": f"c{cid}", "type": "function",
                                     "function": call}]})
        msgs.append({"role": "tool", "tool_call_id": f"c{cid}",
                     "content": result})
    msgs.append({"role": "assistant", "content":
                 "Based on the investigation: the root cause is identified "
                 "and the fix is applied. Summary of changes follows."})
    return tok.apply_chat_template(msgs, tools=TOOLS, tokenize=False)


def make_coding_chat(rng, art, tok):
    p, t = rng.choice(art["code"])
    s = rng.randrange(0, max(1, len(t) - 8000))
    snippet = t[s:s + rng.randint(3000, 9000)]
    ask = rng.choice([
        "Review this code for bugs and suggest improvements:",
        "Explain what this code does and identify any performance issues:",
        "This code crashes intermittently. Find the race condition:",
        "Refactor this for readability and add type hints:",
        "Write unit tests covering the edge cases of this code:",
    ])
    msgs = [
        {"role": "system", "content": "You are an expert code reviewer."},
        {"role": "user", "content": f"{ask}\n\n```python\n{snippet}\n```"},
    ]
    return tok.apply_chat_template(
        msgs, tokenize=False, add_generation_prompt=True)


def _bash_corpus(rng):
    files = (glob.glob(f"{ROOT}/vllm/.buildkite/**/*.sh", recursive=True)
             + glob.glob("/usr/share/bash-completion/completions/*")
             + glob.glob("/etc/init.d/*")
             + glob.glob(f"{ROOT}/vllm/docker/Dockerfile*")
             + glob.glob(f"{ROOT}/vllm/**/*.sh", recursive=True)
             + glob.glob("/etc/profile.d/*.sh"))
    rng.shuffle(files)
    buf, out = "", []
    for p in files:
        t = _read(p)
        if len(t) < 200:
            continue
        buf += f"# ==== {os.path.basename(p)} ====\n{t}\n\n"
        while len(buf) >= MAX_TOKENS * 4:
            out.append(buf[: MAX_TOKENS * 4])
            buf = buf[MAX_TOKENS * 4:]
    if buf:
        out.append(buf)
    return out


def _medqa_chunks(rng, n):
    files = glob.glob("/tmp/medqa/**/*.jsonl", recursive=True)
    rng.shuffle(files)
    out, buf = [], ""
    for f in files:
        for line in open(f, errors="ignore"):
            try:
                buf += json.loads(line)["text"] + "\n\n"
            except (json.JSONDecodeError, KeyError):
                continue
            while len(buf) >= MAX_TOKENS * 4:
                out.append(buf[: MAX_TOKENS * 4])
                buf = buf[MAX_TOKENS * 4:]
                if len(out) >= n:
                    return out
    return out


def build_corpus(tok):
    rng = random.Random(42)
    art = gather_artifacts(rng)
    prompts = []

    # 30% raw code continuation
    buf = ""
    for p in art["code_files"]:
        buf += _read(p) + "\n\n"
        while len(buf) >= MAX_TOKENS * 4 and len(prompts) < 60:
            prompts.append(buf[: MAX_TOKENS * 4])
            buf = buf[MAX_TOKENS * 4:]
        if len(prompts) >= 60:
            break

    # 12% bash / automation scripts
    prompts += _bash_corpus(rng)[:24]

    # 34% agentic tool-calling sessions
    for _ in range(68):
        prompts.append(make_agentic_session(rng, art, tok))

    # 10% coding chat
    for _ in range(20):
        prompts.append(make_coding_chat(rng, art, tok))

    # medical textbook continuation (MedQA corpus, 7%)
    prompts += _medqa_chunks(rng, 14)

    # ~8% general knowledge, medical-weighted (generic clinical knowledge
    # only -- no patient-style data)
    medical = [
        "Explain the differential diagnosis approach for acute chest pain, "
        "covering cardiac, pulmonary, gastrointestinal, and musculoskeletal "
        "causes, and which investigations discriminate between them.",
        "Describe how to systematically read a chest X-ray, common findings "
        "(consolidation, pneumothorax, effusion, cardiomegaly, nodules), and "
        "typical pitfalls in interpretation.",
        "Explain the pharmacology of beta-blockers: mechanism, receptor "
        "selectivity, indications, contraindications, and interactions.",
        "Walk through the pathophysiology of type 2 diabetes from insulin "
        "resistance to complications, and the mechanism of each major drug "
        "class used to treat it.",
        "Explain how CT and MRI imaging work physically, when each is "
        "preferred clinically, and their contraindications.",
        "Describe the interpretation of a full blood count and common "
        "patterns: microcytic vs macrocytic anaemia, neutrophilia, "
        "lymphopenia, thrombocytopenia, and their differential causes.",
        "Explain the staging and grading of solid tumours, TNM notation, "
        "and how imaging and histopathology contribute to each.",
        "Describe the physiology of the cardiac cycle and how it maps to "
        "ECG waveforms, heart sounds, and common arrhythmia mechanisms.",
        "Explain sepsis: definitions, pathophysiology, early recognition "
        "criteria, and the evidence behind initial management bundles.",
        "Describe how radiology AI models are validated for clinical use: "
        "ground truth curation, reader studies, sensitivity/specificity "
        "tradeoffs, and regulatory considerations.",
        "Explain acid-base disturbances and how to interpret an arterial "
        "blood gas step by step, with compensated and mixed examples.",
        "Describe the mechanisms and comparative effectiveness of the major "
        "vaccine platforms: live attenuated, inactivated, subunit, mRNA, "
        "and viral vector.",
    ]
    general = ["the history of the transistor", "the economics of shipping",
               "how compilers optimize loops", "how DNS resolution works",
               "what causes inflation", "the water cycle"]
    for q in medical:
        msgs = [{"role": "user", "content": q + " Be thorough and detailed."}]
        prompts.append(tok.apply_chat_template(
            msgs, tokenize=False, add_generation_prompt=True))
    for t in general:
        msgs = [{"role": "user", "content":
                 f"Give me a long, detailed explanation of {t}, with "
                 "history, fundamentals, examples, and common misconceptions."}]
        prompts.append(tok.apply_chat_template(
            msgs, tokenize=False, add_generation_prompt=True))

    rng.shuffle(prompts)
    return prompts


def main():
    os.environ["VLLM_EXPERT_STATS_DIR"] = STATS_DIR
    from vllm import LLM, SamplingParams
    from vllm.inputs import TokensPrompt

    llm = LLM(
        model=MODEL,
        pipeline_parallel_size=8,
        gpu_memory_utilization=0.6,
        kv_cache_dtype="fp8_ds_mla",
        max_model_len=8192,
        max_num_seqs=4,
        enforce_eager=True,
        kernel_config={"moe_backend": "cutlass"},
    )
    tok = llm.get_tokenizer()
    texts = build_corpus(tok)
    print(f"{len(texts)} calibration prompts", flush=True)
    prompts = [
        TokensPrompt(prompt_token_ids=tok.encode(t)[:MAX_TOKENS])
        for t in texts
    ]
    sp = SamplingParams(max_tokens=1, temperature=0.0)
    outs = llm.generate(prompts, sp)
    total_tokens = sum(len(o.prompt_token_ids) for o in outs)
    del llm

    files = sorted(glob.glob(os.path.join(STATS_DIR, "layer_*.npy")))
    assert files, f"no stats dumped in {STATS_DIR}"
    layer_ids = sorted(int(os.path.basename(f)[6:-4]) for f in files)
    counts = np.stack(
        [np.load(os.path.join(STATS_DIR, f"layer_{li}.npy"))
         for li in layer_ids])
    np.savez(OUT, counts=counts, layer_ids=np.array(layer_ids),
             total_tokens=total_tokens)
    frac = counts / counts.sum(axis=1, keepdims=True).clip(min=1)
    sf = np.sort(frac, axis=1)[:, ::-1]
    print(f"tokens: {total_tokens}, layers: {len(layer_ids)}")
    for k in (16, 32, 64):
        m = sf[:, :k].sum(axis=1)
        print(f"top-{k} mass: mean {m.mean():.3f} min {m.min():.3f} "
              f"max {m.max():.3f}")
    print("saved", OUT)


if __name__ == "__main__":
    main()