prernaaa12/abot-world-interactive

1"""2将 PyTorch CUDA profiler 事件按内核/算子名称粗分为若干类，便于估算3memcpy、GEMM/Linear、Attention、Norm 等在 GPU 时间中的占比。4 5说明：6- 分类基于名称子串启发式，不同 CUDA / Triton / cuBLAS 版本下内核名会有差异；7- 若需更细粒度，请配合 prof.export_chrome_trace() 在 chrome://tracing 或 Perfetto 中查看。8"""9 10from __future__ import annotations11 12from collections import defaultdict13from typing import DefaultDict, Dict, Iterable, List, Tuple14 15# 顺序有意义：先匹配更具体的类别16_BUCKET_ORDER: List[str] = [17    "memcpy_sync",18    "attention",19    "gemm_linear",20    "conv",21    "norm",22    "elementwise_reduce",23    "other",24]25 26 27def classify_cuda_name(name: str) -> str:28    """根据内核或 aten 算子名称归入粗分类。"""29    n = (name or "").lower()30 31    # 设备同步、拷贝、显存设置（含部分 launch 开销）32    if any(33        s in n34        for s in (35            "memcpy",36            "memset",37            "memgetinfo",38            "cudaget",39            "cudaevent",40            "cudaoccupancy",41            "cuda stream",42            "cudastream",43            "cudadevicesynchronize",44            "cuda devicesynchronize",45            "device synchronize",46            "eventrecord",47            "eventsynchronize",48        )49    ):50        return "memcpy_sync"51 52    # FlashAttention / SDPA / 各类 fused attention53    if any(54        s in n55        for s in (56            "flash_attn",57            "flash-attn",58            "scaled_dot_product",59            "sdpa",60            "efficient_attention",61            "mem_eff_attention",62            "fused_attention",63            "fmha",64            "mha_default",65            "multi_head_attention",66            "contrib_attn",67            "attention_forward",68            "attention_backward",69            "softmax",70        )71    ):72        # 独立 softmax 小核也可能与 attention 同桶；若需拆开可把 softmax 挪到 elementwise73        return "attention"74 75    # MatMul / Linear / GEMM（含 cutlass、cublas、triton、fp8）76    if any(77        s in n78        for s in (79            "gemm",80            "cublas",81            "cutlass",82            "matmul",83            "mat_mul",84            "aten::linear",85            "aten::mm",86            "aten::bmm",87            "aten::addmm",88            "aten::matmul",89            "scaled_mm",90            "fp8",91            "wmma",92            "mma_sync",93            "triton",94            "dot",95        )96    ):97        return "gemm_linear"98 99    if any(s in n for s in ("conv", "cudnn", "depthwise", "convolution")):100        return "conv"101 102    if any(103        s in n104        for s in (105            "layernorm",106            "layer_norm",107            "rms_norm",108            "group_norm",109            "flash_norm",110            "aten::layer_norm",111            "aten::group_norm",112            "aten::native_layer_norm",113        )114    ):115        return "norm"116 117    if any(118        s in n119        for s in (120            "elementwise",121            "vectorized",122            "unary",123            "binary",124            "reduce",125            "reduction",126            "activation",127            "silu",128            "gelu",129            "swiglu",130            "relu",131            "aten::add",132            "aten::mul",133            "aten::div",134            "aten::pow",135            "aten::sqrt",136            "aten::rsqrt",137        )138    ):139        return "elementwise_reduce"140 141    return "other"142 143 144def _event_cuda_time_us(event) -> float:145    for attr in (146        "cuda_time_total",147        "self_cuda_time_total",148        "self_cuda_time_total_us",149    ):150        v = getattr(event, attr, None)151        if v is not None:152            return float(v)153    return 0.0154 155 156def aggregate_from_profiler(prof) -> Tuple[Dict[str, float], List[Tuple[str, float]]]:157    """158    从 torch.profiler.profile 实例聚合：159    - 返回 (bucket -> 微秒总和, 按耗时排序的 (name, us) 列表)160    使用 prof.events() 以尽量接近 CUDA 内核级名称。161    """162    bucket_us: DefaultDict[str, float] = defaultdict(float)163    per_name: DefaultDict[str, float] = defaultdict(float)164 165    try:166        events = prof.events()167    except Exception:168        events = []169 170    for e in events:171        us = _event_cuda_time_us(e)172        if us <= 0:173            continue174        name = getattr(e, "name", "") or ""175        b = classify_cuda_name(name)176        bucket_us[b] += us177        per_name[name] += us178 179    # 若 events() 为空，回退到 key_averages（多为 aten 级）180    if not bucket_us and not per_name:181        try:182            for avg in prof.key_averages():183                us = float(getattr(avg, "cuda_time_total", 0) or getattr(avg, "self_cuda_time_total", 0) or 0)184                if us <= 0:185                    continue186                name = getattr(avg, "key", "") or ""187                b = classify_cuda_name(name)188                bucket_us[b] += us189                per_name[name] += us190        except Exception:191            pass192 193    sorted_names = sorted(per_name.items(), key=lambda x: -x[1])194    out = {k: bucket_us.get(k, 0.0) for k in _BUCKET_ORDER}195    for k, v in bucket_us.items():196        if k not in out:197            out[k] = v198    return out, sorted_names199 200 201def format_bucket_report(202    bucket_us: Dict[str, float],203    top_names: Iterable[Tuple[str, float]],204    top_k: int = 25,205) -> str:206    total = sum(bucket_us.values()) or 1e-9207    lines: List[str] = []208    lines.append("=== CUDA 时间按粗分类（微秒 / 占比）===")209    for k in _BUCKET_ORDER:210        if k in bucket_us and bucket_us[k] > 0:211            us = bucket_us[k]212            lines.append(f"  {k:22s}  {us:12.1f} us  ({100.0 * us / total:5.1f}%)")213    # 其它未在顺序表中的桶214    for k, us in sorted(bucket_us.items(), key=lambda x: -x[1]):215        if k in _BUCKET_ORDER:216            continue217        if us > 0:218            lines.append(f"  {k:22s}  {us:12.1f} us  ({100.0 * us / total:5.1f}%)")219    lines.append(f"  {'TOTAL':22s}  {total:12.1f} us")220    lines.append("")221    lines.append(f"=== 耗时最高的 {top_k} 个事件名（微秒）===")222    for i, (name, us) in enumerate(top_names):223        if i >= top_k:224            break225        short = name if len(name) <= 120 else name[:117] + "..."226        lines.append(f"  {us:10.1f}  {short}")227    return "\n".join(lines)228