File size: 2,619 Bytes
fdc6474 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 | #!/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()
|