File size: 1,229 Bytes
0e322d7 49c3379 0e322d7 | 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 | """Per-segment NPU time of one decode step (S=1), plus host overhead.
usage: python prof_segments.py ENGINE_DIR"""
import sys, time
import numpy as np
import openvino as ov
from onw.runtime import SegmentedModel
def main():
core = ov.Core()
m = SegmentedModel(core, sys.argv[1], "NPU", {"CACHE_DIR": "D:/moe-npu/npu_cache"})
m.step(list(range(100, 116)))
times = np.zeros(len(m.segs[1]))
proxies = []
for i, (req, meta, names) in enumerate(m.segs[1]):
class P:
pass
p = P()
p.set_tensor, p.model_inputs = req.set_tensor, req.model_inputs
def g(*a, _f=req.infer, _i=i, **k):
t0 = time.perf_counter()
r = _f(*a, **k)
times[_i] += time.perf_counter() - t0
return r
p.infer = g
proxies.append((p, meta, names))
m.segs[1] = proxies
n = 30
t0 = time.perf_counter()
for i in range(n):
m.step([200 + i])
tot = (time.perf_counter() - t0) / n * 1000
seg = times / n * 1000
print(f"per token {tot:.1f} ms = NPU calls {seg.sum():.1f} ms + host {tot - seg.sum():.1f} ms")
print("segment ms:", " ".join(f"{t:.1f}" for t in seg))
if __name__ == "__main__":
main()
|