| import os, sys, torch, time |
| sys.path.insert(0, "apps"); sys.path.insert(0, "kernels") |
| from common import get_agent; from fast_laya import accelerate |
| from laya.common import build_sequence, collate_items, QTYPES |
| agent = get_agent("multilingual"); fast = accelerate(agent, use_graphs=False) |
| Q = {"department": {"type": "choice", "instructions": "Which team should handle this?", "criteria": {"billing": "invoices, refunds", "technical": "bugs, outages", "sales": "pricing"}}, |
| "urgency": {"type": "score", "instructions": "How urgent is this?", "criteria": ["not urgent", "soon", "blocking"]}, |
| "churn": {"type": "noul", "instructions": "Does the user threaten to cancel?"}} |
| qs = {f"{k}{i}": v for i in range(10) for k, v in Q.items()} |
| for name, st in [("short", {"body": "We were billed twice for March. Please refund the duplicate."}), ("long", {"body": "Since yesterday our whole team cannot log in, the dashboard returns 502 errors. " * 50})]: |
| items = [] |
| for qid in qs: |
| qq = agent._to_internal(qs[qid]); seq, m = build_sequence(agent.tok, st, qq, agent.cfg["max_len"], agent.cfg["head_max_len"]); items.append({"ids": seq, "markers": m, "qtype": QTYPES[qq["t"]]}) |
| b = {k: v.cuda() for k, v in collate_items([items], agent.tok.pad_token_id).items() if torch.is_tensor(v)} |
| def fwd(): |
| with torch.no_grad(), torch.autocast("cuda", dtype=agent.dtype): return agent.model(b["input_ids"], b["attention_mask"], b["marker_pos"], b["marker_mask"], b["qtype"]) |
| for _ in range(3): fwd() |
| from torch.profiler import profile, ProfilerActivity |
| with profile(activities=[ProfilerActivity.CUDA]) as prof: |
| for _ in range(5): fwd() |
| torch.cuda.synchronize() |
| print(f"\n##### {name}: N={b['input_ids'].shape[0]} L0={b['input_ids'].shape[1]}") |
| ka = prof.key_averages() |
| tot = sum(e.self_device_time_total for e in ka) |
| for e in sorted(ka, key=lambda e: -e.self_device_time_total)[:14]: |
| if e.self_device_time_total > 0: print(f"{e.self_device_time_total/5/1000:8.3f} ms {100*e.self_device_time_total/tot:5.1f}% {e.key[:90]}") |
| print(f"{tot/5/1000:8.3f} ms total GPU per forward") |
|
|
