engram_wait.py 886 B

123456789101112
  1. import gzip,json,sys,statistics
  2. for path in sys.argv[1:]:
  3. with gzip.open(path,'rt') as f: data=json.load(f)
  4. kern=sorted([e for e in data['traceEvents'] if e.get('ph')=='X' and e.get('cat')=='kernel'],key=lambda e:e['ts'])
  5. eng=[];other=[]
  6. for i,k in enumerate(kern):
  7. if 'nccl' not in k['name'].lower(): continue
  8. # look back up to 6 kernels for an engram gather
  9. prev=[kern[j]['name'] for j in range(max(0,i-6),i)]
  10. (eng if any('_engram_gather' in n for n in prev) else other).append(k['dur'])
  11. tot=sum(eng)+sum(other)
  12. print(f"{path}: engram-following NCCL: n={len(eng)} total={sum(eng)/1000:.1f} ms mean={statistics.mean(eng) if eng else 0:.0f}us median={statistics.median(eng) if eng else 0:.0f}us max={max(eng) if eng else 0:.0f}us | other NCCL n={len(other)} total={sum(other)/1000:.1f} ms | engram share {100*sum(eng)/tot:.0f}%")