| 12345678910111213141516171819202122232425262728293031323334 |
- #!/usr/bin/env bash
- # Profile N steady-state decode steps on rank 0 (mid-generation, no prefill) and
- # group GPU kernel time by category. Usage: profile_decode.sh LABEL [steps] [prompt]
- set -euo pipefail
- label=${1:?label}; steps=${2:-10}; prompt=${3:-Write a Python function implementing an LRU cache class with get/put, with type hints and docstrings. Code only.}
- dir=$HOME/mimo-v2.6-pro/cache/prof-$label; rm -rf "$dir"
- curl -s localhost:30000/v1/chat/completions -H 'Content-Type: application/json' -d "{\"model\":\"mimo-v2.6-pro\",\"messages\":[{\"role\":\"user\",\"content\":\"$prompt\"}],\"max_tokens\":600,\"temperature\":0,\"chat_template_kwargs\":{\"enable_thinking\":false}}" >/dev/null &
- sleep 2
- curl -s -XPOST localhost:30000/start_profile -H 'Content-Type: application/json' \
- -d "{\"output_dir\":\"/root/.cache/prof-$label\",\"num_steps\":$steps,\"activities\":[\"GPU\"]}" >/dev/null
- wait
- for _ in $(seq 60); do ls "$dir"/*TP-0*.gz >/dev/null 2>&1 && break; sleep 1; done; sleep 3
- python3 - "$dir" "$steps" <<'PY'
- import gzip, json, collections, glob, sys, re
- f = sorted(glob.glob(sys.argv[1] + "/*TP-0*.gz"))[-1]; steps = int(sys.argv[2])
- ev = [e for e in json.load(gzip.open(f))["traceEvents"] if e.get("ph") == "X" and e.get("cat") in ("kernel", "gpu_memcpy", "gpu_memset")]
- span = max(e["ts"] + e["dur"] for e in ev) - min(e["ts"] for e in ev)
- busy = sum(e["dur"] for e in ev)
- rules = [("allreduce-roce", r"roce|Roce"), ("allreduce/gather-nccl", r"nccl"), ("moe-gemm", r"GroupProblemShape|fused_moe|marlin|Marlin"),
- ("moe-misc", r"MoeRouting|expandInputRows|doActivation|ExpertMaps|computeStrides|router|quantize_with_block|topk"),
- ("bf16-gemm (o_proj/lm_head/draft)", r"wmma|nvjet|gemv|splitKreduce|cublas|gemm_bf16"),
- ("fp8-gemm (qkv/dense)", r"w8a8|fp8"), ("attention", r"_fwd_kernel|attn|attention|decode_grouped|flash"),
- ("norm/rope/kv-store", r"rmsnorm|RMSNorm|rotary|kvcache|store_kv"), ]
- cat = collections.defaultdict(float); top = collections.defaultdict(lambda: [0, 0])
- for e in ev:
- n = e["name"]; top[n[:90]][0] += e["dur"]; top[n[:90]][1] += 1
- for c, rx in rules:
- if re.search(rx, n): cat[c] += e["dur"]; break
- else: cat["other"] += e["dur"]
- print("steps %d span %.1f ms (%.1f ms/step) gpu busy %.1f ms (%.0f%%)" % (steps, span/1e3, span/1e3/steps, busy/1e3, 100*busy/span))
- for c, d in sorted(cat.items(), key=lambda x: -x[1]): print(" %-34s %7.2f ms/step %4.1f%%" % (c, d/1e3/steps, 100*d/busy))
- print("top kernels:")
- for n, (d, c) in sorted(top.items(), key=lambda x: -x[1][0])[:12]: print(" %7.2f ms/step %5d %s" % (d/1e3/steps, c, n))
- PY
|
