Triton 内置 Python Profiling API 完全指南：用 `triton.profiler` 精准剖析 Kernel 性能

Triton 内置 Python Profiling API 完全指南用triton.profiler精准剖析 Kernel 性能【免费下载链接】tritonDevelopment repository for the Triton language and compiler项目地址: https://gitcode.com/GitHub_Trending/tri/triton本篇指南系统讲解 Triton 编译器内置性能剖析工具 Proton 的 Python APItriton.profiler。无论你是在训练循环中量化单个 Kernel 的耗时、对比不同配置下的 FLOPS 与带宽还是想为自定义 GPU 算子补充度量指标读完本文你都能熟练使用 Session、Scope、State、Metrics 与 Launch Metadata Hook 构建一套完整、可复现的 Kernel 剖析流程。导入 ProtonProton 通过 Triton 的 profiler 包对外暴露在 Python 中一行即可引入import triton.profiler as proton从源码看triton.profiler包在导入时third_party/proton/proton/init.py会先完成 ROCm/HSA 运行时的选择性配置随后导出start、activate、deactivate、finalize、profile、scope、cpu_timed_scope、enter_scope、exit_scope、state、enter_state、exit_state等全部公开接口。也就是说包导入即完成底层原生库libproton的就绪准备后续所有 API 调用都会经由triton._C.libproton进入 C 实现。Sessions剖析会话的生命周期Session 是 Proton 的顶层抽象代表一段从开始到最终落盘的完整剖析过程。核心 API 为start、activate、deactivate、finalize典型用法如下session proton.start(profile_name, contextshadow) with proton.scope(step): kernelgrid proton.deactivate(session) # Work done here is skipped. proton.activate(session) with proton.scope(next_step): kernelgrid proton.finalize(session)其中deactivate之后到activate之间的代码不会被记录适合跳过预热warmup等不需要统计的阶段。proton.start参数详解proton.start的签名与语义对应 profile.py 中的start实现参数取值说明name任意字符串输出路径不含格式后缀。省略时默认写proton.suffix若datatree实际默认名为~/proton.hatchetcontextshadow/pythonshadow使用 Proton scope 名称组织调用树python使用 Python 文件/函数/行号调用路径datatree/tracetree输出 Hatchet 格式trace输出 Chrome trace 格式backendcupti、rocprofiler、roctracer、instrumentation或NoneNone时根据当前活跃 Triton runtime 自动选择AMD 上优先rocprofilerSDK 不可用时回退roctracermodebackend 特定的模式字符串或模式对象例如 cupti/rocprofiler 支持pcsampling、periodic_flushingroctracer 支持periodic_flushinghooktriton或自定义 hook 实例triton表示启用内核启动元数据记录几个值得注意的源码细节自动选后端_select_backend()通过triton.runtime.driver.active.get_current_target().backend拿到当前目标平台再交给libproton.select_profiler_from_triton_backend映射因此大多数场景无需手动指定backendprofile.py。AMD 环境约束当 backend 为rocprofiler/roctracer时若检测到HIP_VISIBLE_DEVICES或CUDA_VISIBLE_DEVICES被设置会直接抛出ValueError提示改用ROCR_VISIBLE_DEVICESprofile.py。Blackwell 特判当 backend 为cupti且目标架构arch 100Blackwell 及以后时会自动把TRITON_CUPTI_LIB_PATH指向为 Blackwell 构建的 CUPTI 库目录。命令行模式互斥start在脚本被proton命令行工具驱动flags.command_line或剖析被禁用triton.knobs.proton.disable时直接返回None避免重复剖析。mode除了字符串形式还可以传入BaseMode子类实例例如PCSampling(interval1000)会序列化为pcsampling:interval1000字符串见 mode.py。instrumentation后端还支持更细粒度的metric_type如cycle、sampling_strategynone/selective、granularitycta/warp/warp_group及倍数和缓冲策略circular/flush、shared/global等配置。多 Session 管理start返回的 session ID 是整数或剖析关闭时的None。当多个 session 同时活跃时不传 session 的activate()、deactivate()、finalize()会对所有 session生效需要单独控制某个会话时务必把返回的 session 传回对应函数。注意命令行模式下每个函数只能操作一个 sessionsession 为0否则抛出ValueErrorprofile.py。deactivate(session, flushingFalse)支持flushing参数表示在停止记录前是否先冲刷已采集的数据finalize(session, output_format)则负责把数据落盘output_format可选hatchet、hatchet_msgpack、chrome_trace。Function Profiling装饰器一键剖析proton.profile装饰器会在调用被包裹函数前自动start一个 session函数返回后关闭deactivate该 session最后仍需手动调用proton.finalize()落盘proton.profile(nameprofile_name, contextpython) def run(): kernelgrid run() proton.finalize()装饰器也可以不带参数直接使用此时所有配置走默认值contextshadow、datatree、自动选 backendproton.profile def run(): kernelgrid从源码看profile通过_profiling包装函数实现每次调用包裹函数都会start→ 执行函数 →deactivate因此它天然适合每次运行单独出一个剖析会话的脚本形态profile.py。Scopes命名剖析区域Scope 用于在shadow上下文中定义具名区域是组织剖析调用树的基本单元。支持上下文管理器与装饰器两种用法session proton.start(profile_name, contextshadow) with proton.scope(test0): with proton.scope(test1): kernelgrid with proton.scope(test2): kernelgrid proton.finalize(session)作为装饰器proton.scope(matmul) def run_matmul(): matmul_kernelgrid手动 Scope 控制当作用域边界跨越普通with语句不便表达的控制流时可手动配对proton.enter_scope(load) kernelgrid proton.exit_scope(load)exit_scope的参数可以省略但传入名字有助于在作用域不配对时立即暴露错误——源码中exit_scope会把弹出的名字与传入名比较不一致直接抛ValueErrorscope.py。此外手动模式内部用线程局部栈记录未闭合的 scope因此不同线程的 scope 不会互相干扰。Metrics为 Scope 附加性能指标Scope 接受一个 metrics 字典值可以是整数、浮点数、标量张量或整数/浮点数的向量with proton.scope(matmul, {flops16: 2.0e12, bytes: 256_000}): matmul_kernelgrid with proton.scope(candidates, {latency_samples: [10.0, 12.0, 11.0]}): kernelgrid指标后缀语义指标名可携带可选的后缀控制其在调用树中的累积方式后缀含义(inc)或无后缀Inclusive 指标会向上累积到父 scope(exc)Exclusive 指标只保留在当前 scope(pty)Property 指标不跨重复 scope 累积作为属性记录示例with proton.scope(outer, {tokens (inc): 128}): with proton.scope(inner, {cpu_wait (ns)(exc): 90}): kernelgrid proton.enter_scope(config, metrics{tile_size (pty): 128}) proton.exit_scope()指标类型一致性约束同一个指标名在一次剖析中必须使用一致的值类型重复使用相同指标名但传入不兼容的标量类型会报错。从底层实现看指标值在进入 C 层前会经过统一化处理——浮点值转double、整数值转int64向量则对应vector_double/vector_int64类型索引metric.py设备侧张量指标通过两个内置 Triton kerneltensor_metric_kernel与scalar_metric_kernel异步写入设备缓冲区并用原子操作分配记录槽位以支持并发流metric.py。CPU Timed Scopes记录独占 CPU 墙钟时间cpu_timed_scope是scope的子类会在进入/退出时用time.perf_counter_ns()测量独占的 CPU 墙钟时间并写入指标cpu_timewith proton.cpu_timed_scope(host_preprocessing): prepare_inputs()实现细节该 scope 在退出时以{cpu_time (ns)(exc): cpu_time}的形式上报即 CPU 时间是**排他exclusive**的不会向上累积同时cpu_time是保留指标名若你在 metrics 字典里显式传入同名键会直接抛ValueErrorscope.py。随后可在终端用proton-viewer同时对比 GPU 与 CPU 时间proton-viewer -m time/ns,cpu_time/ns profile_name.hatchetStates覆盖 GPU 操作的调用路径标签proton.state用于定制 GPU 操作的调用路径。与 scope 不同state 遵循最近覆盖语义with proton.scope(iteration): with proton.state(optimizer): optimizer_kernelgridStates 与 scopes 的关键区别一个 GPU 操作最多只有一个活跃 state最内层的 state 会覆盖外层 statestate 会被追加到每个 kernel 上方在shadow与python两种上下文中都生效。手动控制同样可用proton.enter_state(state0) kernelgrid proton.exit_state()源码层面state上下文管理器/装饰器直接调用libproton.enter_state/libproton.exit_state并在剖析关闭flags.profiling_on False时安全跳过state.py。内部还提供了一个metadata_state子类用于给内核元数据单独命名空间避免与用户自定义 state 冲突。Launch Metadata Hook自动记录内核启动元数据传入hooktritonProton 会自动记录 Tritonlaunch_metadata回调返回的元数据从而把每个内核的 FLOPS、字节数等指标自动挂到对应的启动节点上from typing import NamedTuple import triton import triton.language as tl import triton.profiler as proton proton.start(profile_name, hooktriton) def metadata_fn(grid: tuple, metadata: NamedTuple, args: dict): return { name: copy_kernel, bytes: args[n_elements] * 4, flops32: 0.0, } triton.jit(launch_metadatametadata_fn) def copy_kernel(x, y, n_elements: tl.constexpr): ...保留元数据键以下键由 Triton runtime 保留用于标识内核与算力指标不会被当作自定义指标name: str flops8: float flops16: float flops32: float flops64: float bytes: int除保留键外的标量或向量数值型元数据会被记录为自定义指标。从实现看LaunchHook._reserved_metadata_keys还额外包含function与streamhooks/launch.py其余Number类型或带data_ptr()的张量值才会被采纳为指标name会作为内核在调用树中的节点名其余值则经transform_tensor_metrics归一化为标量/张量指标后挂到该节点上。自定义LaunchHook与正则过滤进阶用户可以构造并配置LaunchHook实例用正则按内核名过滤需要记录的启动事件from triton.profiler.hooks.launch import LaunchHook hook LaunchHook() hook.configure(include.*matmul.*) proton.start(profile_name, hookhook)include与exclude都是针对编译后内核名的正则表达式仅应使用其中一个过滤方向两者同时配置的语义是先看 include 是否匹配再看 exclude 是否命中见 hooks/launch.py。实现上还做了两段式早退优化内核名能从原始元数据直接判定不匹配时跳过metadata.get()的完整求值从而降低过滤场景下的额外开销。此外LaunchHook是进程级单例多次LaunchHook()构造返回同一实例避免重复注册hook 的进入/退出通过HookManager挂接到 Triton 运行时内核加载与启动的钩子链上kernel_load_end_hook、launch_enter_hook等并按priority排序执行hooks/hook.py。InstrumentationHook会在instrumentation后端下自动注册用于采集指令插桩类指标。命令行快速剖析补充除 Python API 外Proton 还提供等价的命令行入口适合不修改脚本的快速剖析详见 third_party/proton/README.md 与 cli-and-viewer.mdproton [options] script.py [script_args] proton [options] pytest [pytest_args] python -m triton.profiler.proton [options] script.py [script_args]命令行选项与proton.start参数一一对应-n/--name指定会话名、-b/--backend指定后端、-c/--context选择shadow/python、-m/--mode传入模式、-d/--data选择tree/trace、-k/--hook传入triton。命令行模式下脚本内的start调用会被忽略flags.command_line置位保证由命令行统一掌控剖析边界proton.py。小结与推荐工作流一套推荐的端到端剖析流程用proton.start(name, contextshadow, hooktriton)打开会话让 Triton 自动记录内核启动元数据用嵌套with proton.scope(...)组织语义化区域必要时用cpu_timed_scope对照主机侧预处理开销对需要区分执行阶段如 optimizer/forward的调用用proton.state打标签用deactivate/activate跳过预热段最后finalize(session)落盘再用proton-viewer -m time/ns,cpu_time/ns,flops32,bytes profile_name.hatchet在终端直接对比 GPU/CPU 时间与算力指标。上述所有 API 的精确行为都可以在 third_party/proton/proton/profile.py、scope.py、state.py、metric.py 与 hooks/launch.py 中对照阅读配合 cli-and-viewer.md 可进一步掌握视图与分析技巧。【免费下载链接】tritonDevelopment repository for the Triton language and compiler项目地址: https://gitcode.com/GitHub_Trending/tri/triton创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考