#!/usr/bin/env python3
"""Capture SM120 bring-up golden data from the working (SM100) deployment.

Produces in $VLLM_SM120_GOLDEN_DIR (/data/glm52-sm120-golden):
  attn_layer{L}_call{N}.pt  - real sparse-MLA attention I/O incl fp8_ds_mla
                              KV pages and indexer topk selections
  moe_layer{L}_call{N}.pt   - real hybrid-MoE I/O (x, routing, out)
  e2e_goldens.pt            - fixed prompts -> generated ids + top-50
                              logprobs per step (greedy)
"""
import os

import torch

GOLD = "/data/glm52-sm120-golden"
os.environ["VLLM_SM120_GOLDEN_DIR"] = GOLD
os.environ.setdefault("VLLM_PP_LAYER_PARTITION", "21,19,19,19")

PROMPTS = [
    "The capital of France is",
    "def quicksort(arr):\n    ",
    # medium-length: forces a chunked prefill and exercises paged KV
    ("The following is a technical design review.\n\n" +
     "Section {i}: The system shall maintain consistency under partition "
     "by electing a coordinator and journaling all state transitions to "
     "a replicated log with fsync barriers at commit boundaries. " * 220 +
     "\n\nQuestion: What mechanism ensures consistency under partition? "
     "Answer:"),
]


def main():
    from vllm import LLM, SamplingParams

    llm = LLM(
        model="/data/glm52",
        pipeline_parallel_size=4,
        gpu_memory_utilization=0.509,
        kv_cache_dtype="fp8_ds_mla",
        max_model_len=32768,
        max_num_seqs=1,
        max_num_batched_tokens=2048,
        enforce_eager=True,
        max_logprobs=50,
    )
    open(os.path.join(GOLD, 'armed'), 'w').close()
    sp = SamplingParams(max_tokens=32, temperature=0.0, logprobs=50)
    outs = llm.generate(PROMPTS, sp)

    goldens = []
    for o in outs:
        c = o.outputs[0]
        steps = []
        for lp in c.logprobs:
            steps.append({int(t): float(v.logprob) for t, v in lp.items()})
        goldens.append({
            "prompt_token_ids": list(o.prompt_token_ids),
            "generated_token_ids": list(c.token_ids),
            "generated_text": c.text,
            "top50_logprobs_per_step": steps,
        })
    torch.save({"goldens": goldens,
                "config": {"kv_cache_dtype": "fp8_ds_mla",
                           "partition": os.environ["VLLM_PP_LAYER_PARTITION"],
                           "greedy": True, "max_tokens": 32}},
               os.path.join(GOLD, "e2e_goldens.pt"))
    for g in goldens:
        print(f"[{len(g['prompt_token_ids'])} tok] -> {g['generated_text'][:70]!r}")
    print("saved", os.path.join(GOLD, "e2e_goldens.pt"))


if __name__ == "__main__":
    main()