| """Kernel-level view of the SP1 kineto traces: per-step kernel families, per-layer MoE spans and collective pairs. |
| |
| Read only. Step periods follow offload/gb10_budget/trace_budget.py exactly (an iteration's GPU period runs from its first |
| own activity to the next iteration's first own activity; the last iteration has no successor and is dropped), so |
| per-step numbers here line up with SP1's trace-budget.json. All durations are kineto's (MEASURED, profile-affected: |
| SP1 measured decode overhead 1.8-4.5 %, prefill 24.9 %). |
| """ |
|
|
| from __future__ import annotations |
|
|
| import bisect |
| import gzip |
| import re |
| from collections.abc import Iterable, Sequence |
| from dataclasses import dataclass, field |
| from pathlib import Path |
| from typing import IO, cast |
|
|
| from offload.gb10_budget.trace_budget import (CALL_KINDS, GPU_KINDS, MAX_TRACE_BYTES, Rules, bucket, details, events, |
| rules) |
| from offload.prefill_maxlen.evidence import Json, number, record |
|
|
| |
| FAMILIES: tuple[tuple[str, re.Pattern[str]], ...] = tuple((f, re.compile(p)) for f, p in ( |
| ('l2pf', r'\bl2pf_kernel\b'), |
| ('smoe_rot', r'\bsmoe_rot\b'), |
| ('smoe_gu', r'\bsmoe_gu\b'), |
| ('smoe_down', r'\bsmoe_down\b'), |
| ('smoe_fused', r'\bsmoe_fused\b'), |
| ('router', r'\brouter_kernel\b'), |
| ('route_topk', r'\b_dsv4_topk_kernel\b'), |
| ('tf_gu', r'\bgrouped_mma\w*kernel\b'), |
| ('tf_down', r'\bgrouped_rows\w*kernel\b'), |
| ('tf_aux', r'\brot_in_kernel\b|\bgateup_epilogue_kernel\b|\bdown_combine_kernel\b|\bdecode_prep_kernel\b|' |
| r'\bwork_list_kernel\b|\bgroup_(count|place)?_?kernel\b|\bcombine(_y)?_kernel\b'), |
| ('mrow_head', r'\bmrow(32)?_kernel<6\b'), |
| ('mrow', r'\bmrow(32)?_kernel<'), |
| ('exl3_gemm', r'\bexl3_gemm_kernel<'), |
| ('reconstruct', r'\breconstruct_had_kernel<'), |
| ('marlin', r'[Mm]arlin'), |
| ('hc_prenorm', r'sm120_tf32_hc_prenorm_gemm'), |
| ('mqa_logits', r'paged_mqa_logits'), |
| ('gemm_other', r'cutlass|nvjet|cublas|gemm|Gemm|GEMM|xmma'), |
| ('coll_stage', r'\bstage_kernel<'), |
| ('coll_sum', r'\bsum_kernel\('), |
| ('coll_gather', r'::gather_kernel<'), |
| ('nccl', r'nccl'), |
| ('sparse_mla', r'sparse_mla'), |
| )) |
| COLLECTIVE = ('coll_stage', 'coll_sum', 'coll_gather', 'nccl') |
| TEMPLATE_K = re.compile(r'(?:mrow(?:32)?_kernel|exl3_gemm_kernel|reconstruct_had_kernel)<(\d+)') |
|
|
|
|
| def family(name: str) -> str: |
| return next((f for f, p in FAMILIES if p.search(name)), 'rest') |
|
|
|
|
| def template_k(name: str) -> int | None: |
| m = TEMPLATE_K.search(name) |
| return int(m.group(1)) if m else None |
|
|
|
|
| @dataclass(frozen=True) |
| class Kernel: |
| name: str |
| family: str |
| bucket: str |
| start: float |
| end: float |
| stream: int |
| grid: tuple[int, int, int] |
| block: tuple[int, int, int] |
| blocks_per_sm: float | None |
| corr: int | None |
|
|
| @property |
| def dur(self) -> float: |
| return self.end - self.start |
|
|
|
|
| @dataclass(frozen=True) |
| class Call: |
| name: str |
| start: float |
| end: float |
| tid: str |
| corr: int | None |
|
|
|
|
| @dataclass(frozen=True) |
| class Annotation: |
| name: str |
| start: float |
| end: float |
| tid: str |
|
|
|
|
| @dataclass |
| class RankTrace: |
| kernels: list[Kernel] |
| calls: list[Call] |
| annotations: list[Annotation] |
|
|
|
|
| @dataclass |
| class Step: |
| index: int |
| annotation: str |
| phase: str |
| rows: int |
| requests: int |
| lo: float |
| hi: float |
| kernels: list[Kernel] = field(default_factory=list) |
|
|
| @property |
| def period(self) -> float: |
| return self.hi - self.lo |
|
|
|
|
| def _triple(value: Json) -> tuple[int, int, int]: |
| if not isinstance(value, list) or len(value) != 3: |
| return (0, 0, 0) |
| a, b, c = (int(number(v)) for v in value) |
| return (a, b, c) |
|
|
|
|
| def parse(handle: IO[str], r: Rules) -> RankTrace: |
| kernels: list[Kernel] = [] |
| calls: list[Call] = [] |
| annotations: list[Annotation] = [] |
| for event in events(handle): |
| cat = event.get('cat') |
| if event.get('ph') != 'X' or not isinstance(cat, str): |
| continue |
| name = str(event.get('name', '')) |
| start = number(event['ts']) |
| end = start + number(event.get('dur', 0)) |
| args = record(event.get('args') or {}) |
| corr = None if args.get('correlation') is None else int(number(args['correlation'])) |
| if cat in GPU_KINDS: |
| occ = args.get('occupancy') |
| bps = None |
| if isinstance(occ, dict) and occ.get('activeBlocksPerMultiprocessor') is not None: |
| bps = number(occ['activeBlocksPerMultiprocessor']) |
| stream = args.get('stream') |
| kernels.append(Kernel(name, family(name), bucket(name, r), start, end, |
| -1 if stream is None else int(number(stream)), _triple(args.get('grid')), |
| _triple(args.get('block')), bps, corr)) |
| elif cat in CALL_KINDS: |
| calls.append(Call(name, start, end, str(event.get('tid')), corr)) |
| elif cat == 'user_annotation' and name.startswith('execute_'): |
| annotations.append(Annotation(name, start, end, str(event.get('tid')))) |
| kernels.sort(key=lambda k: (k.start, k.end)) |
| calls.sort(key=lambda c: c.start) |
| annotations.sort(key=lambda a: a.start) |
| return RankTrace(kernels, calls, annotations) |
|
|
|
|
| def load(path: Path, r: Rules | None = None) -> RankTrace: |
| if path.stat().st_size > MAX_TRACE_BYTES: |
| raise ValueError(f'{path.name}: larger than the per-file cap') |
| with gzip.open(path, 'rt', encoding='utf-8') as handle: |
| trace = parse(cast(IO[str], handle), r or rules()) |
| while handle.read(1 << 20): |
| pass |
| return trace |
|
|
|
|
| def steps(trace: RankTrace) -> list[Step]: |
| """Complete iterations with their GPU period [first own activity, next iteration's first own activity).""" |
| by_corr: dict[int, float] = {} |
| for k in trace.kernels: |
| if k.corr is not None: |
| by_corr[k.corr] = min(by_corr.get(k.corr, k.start), k.start) |
| call_starts = [c.start for c in trace.calls] |
| firsts: list[float | None] = [] |
| for ann in trace.annotations: |
| lo: float | None = None |
| for c in trace.calls[bisect.bisect_left(call_starts, ann.start):bisect.bisect_right(call_starts, ann.end)]: |
| if c.tid == ann.tid and c.corr is not None and c.corr in by_corr: |
| t = by_corr[c.corr] |
| lo = t if lo is None else min(lo, t) |
| firsts.append(lo) |
| starts = [k.start for k in trace.kernels] |
| out: list[Step] = [] |
| for i, ann in enumerate(trace.annotations[:-1]): |
| lo, hi = firsts[i], firsts[i + 1] |
| if lo is None or hi is None or hi <= lo: |
| continue |
| d = details(ann.name) |
| phase = str(d['phase']) |
| rows = d['generation_tokens'] if phase == 'decode' else d['context_tokens'] |
| reqs = d['generation_requests'] if phase == 'decode' else d['context_requests'] |
| mine = trace.kernels[bisect.bisect_left(starts, lo):bisect.bisect_left(starts, hi)] |
| out.append(Step(i, ann.name, phase, int(number(rows or 0)), int(number(reqs or 0)), lo, hi, list(mine))) |
| return out |
|
|
|
|
| def union(spans: Iterable[tuple[float, float]], lo: float = float('-inf'), hi: float = float('inf')) -> float: |
| total = 0.0 |
| cur_s = cur_e = None |
| for s, e in sorted((max(s, lo), min(e, hi)) for s, e in spans): |
| if e <= s: |
| continue |
| if cur_e is None or s > cur_e: |
| if cur_e is not None and cur_s is not None: |
| total += cur_e - cur_s |
| cur_s, cur_e = s, e |
| else: |
| cur_e = max(cur_e, e) |
| if cur_e is not None and cur_s is not None: |
| total += cur_e - cur_s |
| return total |
|
|
|
|
| @dataclass(frozen=True) |
| class MoeLayer: |
| """One layer's routed-expert launches in a step (decode: smoe rot/gu/down; prefill: TF grouped gate|up / down).""" |
| router: Kernel | None |
| topk: Kernel | None |
| rot: Kernel | None |
| gu: Kernel |
| down: Kernel |
|
|
| @property |
| def start(self) -> float: |
| return self.rot.start if self.rot is not None else self.gu.start |
|
|
| @property |
| def span(self) -> float: |
| return max(self.down.end, self.gu.end) - self.start |
|
|
| @property |
| def down_after_gu(self) -> float: |
| return max(0.0, self.down.end - self.gu.end) |
|
|
|
|
| def moe_layers(step: Step) -> list[MoeLayer]: |
| ks = step.kernels |
| out: list[MoeLayer] = [] |
| router: Kernel | None = None |
| topk: Kernel | None = None |
| rot: Kernel | None = None |
| gu: Kernel | None = None |
| for k in ks: |
| if k.family == 'router': |
| router = k |
| elif k.family == 'route_topk': |
| topk = k |
| elif k.family in ('smoe_rot',): |
| rot = k |
| elif k.family in ('smoe_gu', 'tf_gu'): |
| gu = k |
| elif k.family in ('smoe_down', 'tf_down') and gu is not None: |
| out.append(MoeLayer(router, topk, rot, gu, k)) |
| router = topk = rot = gu = None |
| return out |
|
|
|
|
| def family_totals(step: Step) -> dict[str, dict[str, float]]: |
| """Per family: launches, summed kernel time and union (wall) time inside the step period.""" |
| groups: dict[str, list[Kernel]] = {} |
| for k in step.kernels: |
| groups.setdefault(k.family, []).append(k) |
| return {f: {'n': float(len(v)), 'sum_us': sum(min(k.end, step.hi) - k.start for k in v), |
| 'union_us': union(((k.start, k.end) for k in v), step.lo, step.hi)} |
| for f, v in sorted(groups.items())} |
|
|
|
|
| def collective_calls(step: Step) -> list[Kernel]: |
| return [k for k in step.kernels if k.family in COLLECTIVE] |
|
|
|
|
| def compute_segments(step: Step) -> list[tuple[str, float]]: |
| """Main-stream work between collectives: (kind, us) from the end of one all-reduce's sum to the next stage. Kind is |
| 'moe' when the segment holds smoe/TF expert kernels, 'attn' otherwise (the attention half of a layer).""" |
| marks = [k for k in step.kernels if k.family in ('coll_stage', 'coll_sum')] |
| segs: list[tuple[str, float]] = [] |
| prev_end: float | None = None |
| for k in marks: |
| if k.family == 'coll_stage' and prev_end is not None: |
| inside = [x for x in step.kernels if prev_end <= x.start < k.start] |
| kind = 'moe' if any(x.family in ('smoe_gu', 'tf_gu') for x in inside) else 'attn' |
| segs.append((kind, k.start - prev_end)) |
| prev_end = None |
| elif k.family == 'coll_sum': |
| prev_end = k.end |
| return segs |
|
|
|
|
| def pair_steps(a: Sequence[Step], b: Sequence[Step]) -> list[tuple[Step, Step]]: |
| """Steps of the two ranks with the same iteration index and annotation (the ranks run the same schedule).""" |
| bi = {s.index: s for s in b} |
| return [(s, bi[s.index]) for s in a if s.index in bi and bi[s.index].annotation == s.annotation] |
|
|
