TileLang IKET 工具详解为 CUDA Kernel 注入事件标记并生成 Perfetto 性能剖析轨迹【免费下载链接】tilelangDomain-specific language designed to streamline the development of high-performance GPU/CPU/Accelerators kernels项目地址: https://gitcode.com/GitHub_Trending/ti/tilelangIKETIKET Profiling是 TileLang 针对 CUDA 后端提供的一套实验性仪器化剖析工具它通过在tilelang.compile(...)生成 CUDA 源码时注入命名标记marker、warp 级作用域range以及可选的 32 位标量 payload再由外部 IKET profiler 收集事件并导出.pftrace/.trace.json/.html等产物最终可在 Perfetto 中检视。本文基于仓库文档 docs/tools/iket.md 及tilelang/tools/cuda/iket/下的实现源码系统讲解 IKET 的环境要求、事件 API、payload 机制、编译会话与缓存行为、轨迹采集与查看方法以及底层仪器化如何穿越编译流水线的实现原理帮助读者为自己的 TileLang CUDA kernel 建立可观测性。1. 工具定位与使用前提IKET 是一个CUDA 工具不属于 TileLang 语言命名空间tilelang.language。正确的导入方式是import tilelang.language as T from tilelang.tools.cuda import iket这一点对应的实现位于 tilelang/tools/cuda/iket/init.py该模块将前端事件 APIfrontend、会话生命周期session、宿主侧输出辅助cli统一导出公开符号包括mark、range、range_push/range_pop、payload、session、enable/disable、profile_command、trace_files等。IKET 的接入完全走 TileLang 常规的targetcuda后端不依赖 TileScale 或 CuTe DSL 前端它利用 TileLang 的 CUDA 源码后处理回调tilelang_callback_cuda_postproc在生成 CUDA 文本后注入仪器化代码回调注册逻辑见 tilelang/tools/cuda/iket/session.py 的enable()函数回调实现入口为 tilelang/tools/cuda/iket/codegen.py 的inject_iket_cuda(...)。环境要求目标环境需要满足带 CUDA 支持的 TileLangCUDA 可用的 GPU 及驱动运行仓库自带示例时还需要 PyTorch外部的 IKET Python 包与运行时用于采集轨迹。可以用一条命令验证当前 Python 环境python -c import tilelang, iket, torchSM 架构适配TileLang 的代码生成路径同时支持 Hopper 之前与 Hopper 及以后的目标。从源码看codegen.py 中的_cluster_rank_instruction(...)会根据 target 解析出的 SM 版本决定 rank 指令SM90 及更新架构生成的仪器化 PTX 内联汇编读取%cluster_ctarankmov.b32 r, %cluster_ctarank;Hopper 之前的目标或无法判断目标架构时使用 cluster rank0mov.u32 r, 0;不会发射 Hopper 专属的寄存器。这一细节决定了 IKET 在 A100 与 H100 等平台上都能生成合法的汇编但集群维度行为不同。2. 快速上手在 iket.session 中构建并编译最小可用流程是在iket.session(...)上下文内构建并编译带仪器化的 kernelimport tilelang import tilelang.language as T from tilelang.tools.cuda import iket def instrumented_add(n: int, threads: int 128): T.prim_func def main( A: T.Tensor((n,), T.float32), B: T.Tensor((n,), T.float32), C: T.Tensor((n,), T.float32), ): with T.Kernel(T.ceildiv(n, threads), threadsthreads) as bx: with iket.range(block_total): for tx in T.Parallel(threads): i bx * threads tx if i n: iket.mark(before_store) C[i] A[i] B[i] iket.mark(after_store) return main with iket.session(output_dir/tmp/tilelang_iket): program instrumented_add(1024) kernel tilelang.compile( program, out_idx-1, targetcuda, execution_backendcython, )关键约束与细节会话必须在tilelang.compile(...)生成 CUDA 源码期间处于激活状态否则回调不会被触发生成的 CUDA 源码里不会有 IKET 元数据。推荐在会话内构建 kernel因为此时会话会拿到一份干净的事件注册表reset_eventsTrue会调用前端reset()清空事件 ID 分配但这不是必需的——事件元数据以 token 形式内嵌在 TIR 中因此进入会话之前构建好的PrimFunc在编译时仍携带全部所需信息原理见第 6 节。直接运行程序如仓库示例 examples/iket/minimal.py、examples/iket/all_features.py会编译并执行仪器化 kernel同时通过kernel.get_kernel_source()校验源码中是否包含__iket_meta_info与TL_IKET_EVENT符号并写出.cu源文件。要真正采集轨迹必须通过第 5 节的外部 IKET profiler 运行。仓库提供了三个由浅入深的可运行示例其覆盖范围在 examples/iket/README.md 中有说明示例文件覆盖内容examples/iket/minimal.py瞬时标记marker与一个 warp 局部 rangeexamples/iket/payload_minimal.py一个int32运行时 payload 标记examples/iket/all_features.pyrange、无 payload 标记、运行时 payload 标记的综合演示直接运行仅验证编译与执行不采集轨迹python examples/iket/minimal.py \ --iket-output-dir /tmp/tilelang_iket_minimal python examples/iket/payload_minimal.py \ --iket-output-dir /tmp/tilelang_iket_payload_minimal \ --iket-runtime-payloads python examples/iket/all_features.py \ --iket-output-dir /tmp/tilelang_iket_all_features \ --iket-runtime-payloads--iket-runtime-payloads是一个布尔开关示例内部会将其传给iket.session(runtime_payloads...)控制元数据是否以非 NoPayload形式发射见第 4 节。3. 事件 APIMarkers 与 Ranges3.1 三种基本用法瞬时事件用iket.mark(...)iket.mark(load_inputs)作用域用iket.range(...)作为 Python 上下文管理器包裹一段词法区域with iket.range(compute): # TileLang statements ...显式 push/pop API同样可用iket.range_push(compute) # TileLang statements iket.range_pop(compute)其中iket.range_start(...)与iket.range_end(...)分别是range_push(...)与range_pop(...)的别名见 frontend.py 第 109–116 行。range start 可以携带 payload但range-end 事件不能。3.2 warp 粒度与命名约束warp 粒度记录IKET 以 warp 为粒度记录 range。一个包含 4 个 warp 的 block对一个词法iket.range(...)作用域会产出4 条轨迹 range。名称长度事件与 range 名称上限为32 个 UTF-8 字节。源码中该限制由 metadata.py 的MAX_EVENT_NAME_BYTES 32定义前端_get_event(...)在注册时即校验超出直接抛ValueError。payload 类型冲突拒绝在同一个前端注册表中用不同的 payload dtype 复用同一 marker/range 名称会被拒绝。实现上frontend.py 的_validate_payload_compat(...)会比较首次注册的payload_type与当前使用的类型不一致即抛错。range 的 ID 派生range_id 取名称的 CRC32zlib.crc32因此同名的 range 跨 kernel 也具有一致的 range 标识而event_id由前端注册表顺序分配并且实现刻意跳过 31 这个 ID_RANGE_END_EVENT_ID 31保留给 range-end 专用事件见 frontend.py 第 16 行。push/pop 配对检查range_pop(...)使用线程局部栈校验配对pop 空栈或名称不匹配都会抛RuntimeError这有助于尽早发现不匹配的作用域书写错误。4. 运行时 Payload 机制4.1 支持的类型与描述符marker 和 range start 可以捕获一个 32 位标量值。TileLang 目前支持的 payload dtype 只有三种int32uint32float32对 TileLang 表达式应使用显式 payload 描述符iket.mark(store_index, payloadiket.payload(i, dtypeint32)) iket.mark(scale, payloadiket.payload(value, dtypefloat32))iket.payload(...)返回一个PayloadSpec表达式 dtype IKET 侧类型 ID的冻结 dataclass。简单 Python 标量int/float/bool和带dtype属性的表达式可以直接传入dtype 会被自动推断bool → uint32、int → int32、float → float32但显式指定 dtype 能让 trace 的 schema 无歧义。dtype 别名int/uint/float会被归一化到int32/uint32/float32其它类型直接触发NotImplementedError。4.2 运行时捕获是 opt-inwith iket.session(runtime_payloadsTrue): program instrumented_add(1024) kernel tilelang.compile(program, targetcuda)不开启runtime_payloadsTrue时的行为差异由 codegen.py 的_metadata_decls(...)与_event_macros(...)决定未开启payload schema 仍然编码在 TIR 元数据 token 里但发射的 IKET 元数据声明NoPayload生成的 payload 宏TL_IKET_EVENT_PAYLOAD_U32/F32退化为普通TL_IKET_EVENT(ID)即不写 payload 值。这样普通 marker 记录保持在 4 字节。已开启payload 事件会发射两条独立的 32 位 shared-memory store——一条写时间戳|事件 ID另一条写 payload 值详见第 6.3 节的宏实现。payload 值通过 IKET 的warp 级 dump 机制观察一个 payload 值通常代表该机制选中的 lane而不是 warp 内每个线程都上报同一个值。5. 采集与查看轨迹5.1 用外部 profiler 采集轨迹外部 profiler 通过 IKET 包的 CLI 暴露python -m iket.cli.main对综合示例 examples/iket/all_features.py 的完整采集命令rm -rf /tmp/tilelang_iket_all_features_profile python -m iket.cli.main \ --output-dir /tmp/tilelang_iket_all_features_profile \ --clobber \ profile \ --postprocess all \ -- \ python examples/iket/all_features.py \ --iket-output-dir /tmp/tilelang_iket_all_features_profile \ --iket-runtime-payloadsprofiler 会配置外部 IKET 运行时并在--之后启动目标命令。输出目录中可能包含iket_pid_0x....pftrace iket_pid_0x....pftrace.gz iket_pid_0x....trace.json iket_pid_0x....htmlTileLang 还可以从 Python 侧构造同一条 shell 命令command iket.profile_command( [python, examples/iket/all_features.py, --iket-runtime-payloads], directory/tmp/tilelang_iket_all_features_profile, ) print(command)profile_command(...)只返回带引号转义的命令字符串实现见 cli.py内部按python -m iket.cli.main --output-dir dir --clobber profile --postprocess all -- command拼装并用shlex.quote转义它本身不会启动 profiler。5.2 在 Perfetto 中查看生成的 HTML 需要能加载同目录下的 trace 文件因此应把输出目录作为静态站点服务cd /tmp/tilelang_iket_all_features_profile python3 -m http.server 8080然后在浏览器打开具体生成的文件例如http://localhost:8080/iket_pid_0x....html远程主机上用 SSH 端口转发ssh -L 8080:localhost:8080 userremote-host如果页面只显示 Perfetto 的落地页而没有数据请在 Perfetto UI 中手动导入对应的.pftrace文件。5.3 编程式解析 JSON 导出JSON 轨迹可以程序化检查例如提取所有store_indexmarker 的 payload 值import json from pathlib import Path trace_path max( Path(/tmp/tilelang_iket_all_features_profile).glob(*.trace.json), keylambda path: path.stat().st_size, ) data json.loads(trace_path.read_text()) launch data[launches][0] names data[stringTable] store_indices [ marker[payloadVal] for marker in launch[markers] if names[marker[markerNameIdx]] store_index and payloadVal in marker ] print(store_indices[:8])examples/iket/minimal.py 里也展示了如何用iket.trace_files(...)找到最新轨迹并统计launches/markers/ranges数量。6. 实现原理仪器化如何穿越编译6.1 会话参数与状态回滚iket.session(...)的完整签名with iket.session( reset_eventsTrue, overrideTrue, disable_on_exitTrue, output_dirNone, runtime_payloadsNone, disable_cacheTrue, ): ...各参数控制的宿主侧状态如下参数默认作用reset_eventsTrue清空此后构建 kernel 的前端事件分配已内嵌在PrimFunc里的元数据不受影响overrideTrue允许 IKET 在最外层会话激活期间替换已存在的tilelang_callback_cuda_postproc回调disable_on_exitTrue退出时恢复之前的回调限定作用域使用请保持默认output_dirNone创建目录、设置TL_IKET_OUTPUT_DIR环境变量并在会话期间配置 TileLang 的 IKET 路径辅助函数runtime_payloadsNone临时选择是否发射 payload 值None保持原有设置disable_cacheTrue默认绕过 TileLang 的KernelCache会话退出时恢复其先前状态从 session.py 的实现看_Session.__enter__会先对输出目录、payload 模式、缓存状态做快照再依次应用新状态并注册回调任何一步失败都会走except分支回滚已生效的部分。__exit__则按disable_on_exit恢复回调并回滚快照——这正是优先使用iket.session(...)的原因即使编译抛异常清理也一定发生。为什么disable_cache默认开启因为事件名与 schema 虽然是 TIR 缓存标识的一部分结构相同但事件名不同的 kernel 会得到不同的缓存 key但回调激活与 payload 模式是宿主侧编译状态。复用一份未带回调、或以不同 payload 模式编译的二进制会产生缺失或过期的仪器化。只有在调用方能控制这些条件时才应设disable_cacheFalse。回调是引用计数的enable()/disable()维护_enable_depth嵌套 IKET 会话保持外层回调激活退出最外层会话时才恢复 IKET 之前注册的回调输出目录、payload 模式、缓存状态同样在会话退出后恢复。更底层的手动生命周期 API 如下用于高级场景但需自行保证清理iket.enable() iket.is_enabled() iket.disable() iket.enable_runtime_payloads() iket.runtime_payloads_enabled() iket.disable_runtime_payloads()6.2 元数据 token事件信息内嵌于 TIR前端每次事件调用都会生成一个规范化的元数据 token并作为参数写入 TIR 中的call_extern调用。token 的构造见 metadata.py前缀固定为__tl_iket_v1_内容是{version, name, event_id, kind, range_id, payload_type, payload_iket_id}的 JSON经sort_keys压缩后做 URL-safe base64 编码解码端decode_event(...)对前缀、字段集合、版本、名称长度、kind 取值仅mark/range、ID 合法性event_id不能为 0 或 31逐项校验防止陈旧或损坏的元数据进入代码生成。这带来两个重要后果预构建的PrimFunc在会话进入与前端注册表重置之后仍保留事件元数据——事件信息不依赖进程级注册表结构相似但事件名不同的 kernel 具有不同的 IR 缓存标识——IKET 事件参与 TileLangKernelCache的 key 计算因此改事件名会触发重新编译而不是命中旧二进制。6.3 CUDA 源码后处理元数据数组与事件宏会话激活期间tilelang_callback_cuda_postproc触发 codegen.py 的inject_iket_cuda(...)完成四件事恢复事件用正则_EVENT_CALL_PATTERN从生成的 CUDA 中找出TL_IKET_EVENT_PAYLOAD_U32|_F32调用并解码 tokentoken 数量与事件调用数量必须一致否则报IKET metadata token is not attached to a supported event call。分配模块级事件 ID按(kind, name)归并重新分配从 1 开始的模块级 ID同样跳过 31冲突的元数据会直接抛错。发射 IKET 元数据数组每个事件生成一个 60 字节__device__常量数组事件 ID、instrument method、payload 类型 ID、range 信息、名称等range 另生成 72 字节数组含 CRC32 range_id外加一个 48 字节的__iket_meta_info头含 IKET magic 字节157, 241, 190, 186、条目数、事件上限等。这些数组以extern C声明、__attribute__((used, aligned(1)))保证不被优化掉是外部 IKET patcher 定位与解析事件的依据。定义 NativeDump 事件宏核心宏TL_IKET_EVENT(ID, ...)展开为一段 PTX 内联汇编.reg .b32 r, t; mov.b32 r, %cluster_ctarank; // 或 mov.u32 r, 0; mov.u32 t, %globaltimer_lo; or.b32 t, t, ID; mad.lo.u32 r, r, 0x1000000, 0x20; st.weak.shared.u32 [r], t; pmevent.mask ID;即以 cluster rank 计算 shared memory 槽位每 rank 预留0x1000000字节空间、偏移0x20起把globaltimer 低 32 位 | 事件 ID写入该槽位再触发pmevent.mask性能监控事件。无 payload 事件只写这一条 32 位记录。开启runtime_payloads后payload 宏改为两段独立汇编第一段写时间戳记录偏移0x20第二段把 payload 值以st.volatile.shared.u32写到偏移0x24最后触发pmevent.mask。payload store 特意使用volatile目的是阻止 ptxas 把这对 store 合并为STS.64——该 64 位记录形态是当前外部 IKET patcher不接受的。这是理解第 8 节故障排查的关键。6.4 宿主侧辅助函数CUDA 工具还提供了小型宿主辅助 API实现见 cli.pyiket.set_output_dir(/tmp/tilelang_iket) # 创建目录并设置 TL_IKET_OUTPUT_DIR iket.output_dir() # 返回当前输出目录 iket.output_path(kernel.cu) # 输出目录下的具体文件路径 iket.trace_files() # 返回 *.trace.json按大小降序 iket.profile_command([...], directory/tmp/tilelang_iket) # 只构造命令字符串这些辅助只管理路径与命令构造本身不采集轨迹。set_output_dir(...)会同时写进程变量与环境变量TL_IKET_OUTPUT_DIR示例脚本的--iket-output-dir默认值也正是回落到该环境变量。iket.event_table()返回最近构建 kernel 期间注册的事件列表name、event_id、kind、range_id、payload_type 等适合检查注册结果但不是代码生成的权威来源——权威来源是 TIR 中的 token。若需要显式重置前端事件分配调用iket.reset()。7. 限制Limitations文档明确列出当前 IKET 集成的边界使用与评估时应以此为准仅支持 TileLang 的CUDA 后端运行时 payload 仅限int32、uint32、float32事件与 range 名称上限32 UTF-8 字节不生成源码位置表trace 中的locIdx是 IKET 运行时的位置索引不是Python 或 TIR 行号IKET 以warp 粒度记录事件回调、payload 模式、事件注册表、kernel-cache 开关都是进程全局状态并行的编译工作流需要自行协调访问仓库测试 testing/python/tools/test_tilelang_tools_cuda_iket.py 的autousefixture 就专门做这种状态隔离与恢复;该集成依赖外部 IKET 运行时的私有元数据与 NativeDump 约定应视为实验性能力。8. 故障排查8.1 payload schema 存在但 trace 里没有payloadValschema 已内嵌但值未发射通常是没有在正确的 payload 模式下编译with iket.session(runtime_payloadsTrue): kernel tilelang.compile(program, targetcuda)同时检查该 marker 是否带有受支持的 payload 描述符iket.payload(...)且 dtype 属于 int32/uint32/float32。注意disable_cache默认开启就是为了避免命中未开 payload 模式时编译的旧二进制。8.2 profiler 在给 payload kernel 打补丁时失败payload 仪器化依赖两条独立的 32 位 store。可用反汇编检查生成二进制的记录形态nvdisasm kernel.cubin | grep -E STS|PMTRIG期望的形态是STS [addr], timestamp_with_event_id STS [addr0x4], payload_value PMTRIG event_id如果出现STS.64的时间戳/payload 对说明仪器化序列不再匹配 IKET 的 NativeDump 打补丁约定——检查编译时是否处于runtime_payloads会话、事件宏是否被后续处理改动以及外部 IKET 运行时版本是否与 TileLang 侧约定的宏实现兼容。9. 小结与延伸阅读IKET 把kernel 内的事件观测点做成了 TileLang 编译流水线的一等公民前端在 TIR 中内嵌自描述的元数据 token编译期后处理恢复 token、生成设备侧元数据数组与 PTX 事件宏外部 IKET 运行时再完成采样与轨迹导出。对需要剖析 TileLang CUDA kernel 内部阶段耗时例如 flash attention 中 load/compute/store 各阶段、block 级循环进度、关键索引值的开发者这套机制提供了比整体 kernel 计时更细的视角同时通过会话/引用计数/状态快照保证了与宿主编译状态的隔离与恢复。完整指南文档docs/tools/iket.md工具实现tilelang/tools/cuda/iket/init.py、tilelang/tools/cuda/iket/frontend.py、tilelang/tools/cuda/iket/session.py、tilelang/tools/cuda/iket/codegen.py、tilelang/tools/cuda/iket/metadata.py、tilelang/tools/cuda/iket/cli.py可运行示例examples/iket/README.md、examples/iket/minimal.py、examples/iket/payload_minimal.py、examples/iket/all_features.py行为测试回调注册/恢复、缓存、payload 元数据、SM90 rank 指令等testing/python/tools/test_tilelang_tools_cuda_iket.py【免费下载链接】tilelangDomain-specific language designed to streamline the development of high-performance GPU/CPU/Accelerators kernels项目地址: https://gitcode.com/GitHub_Trending/ti/tilelang创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考
TileLang IKET 工具详解：为 CUDA Kernel 注入事件标记并生成 Perfetto 性能剖析轨迹
TileLang IKET 工具详解为 CUDA Kernel 注入事件标记并生成 Perfetto 性能剖析轨迹【免费下载链接】tilelangDomain-specific language designed to streamline the development of high-performance GPU/CPU/Accelerators kernels项目地址: https://gitcode.com/GitHub_Trending/ti/tilelangIKETIKET Profiling是 TileLang 针对 CUDA 后端提供的一套实验性仪器化剖析工具它通过在tilelang.compile(...)生成 CUDA 源码时注入命名标记marker、warp 级作用域range以及可选的 32 位标量 payload再由外部 IKET profiler 收集事件并导出.pftrace/.trace.json/.html等产物最终可在 Perfetto 中检视。本文基于仓库文档 docs/tools/iket.md 及tilelang/tools/cuda/iket/下的实现源码系统讲解 IKET 的环境要求、事件 API、payload 机制、编译会话与缓存行为、轨迹采集与查看方法以及底层仪器化如何穿越编译流水线的实现原理帮助读者为自己的 TileLang CUDA kernel 建立可观测性。1. 工具定位与使用前提IKET 是一个CUDA 工具不属于 TileLang 语言命名空间tilelang.language。正确的导入方式是import tilelang.language as T from tilelang.tools.cuda import iket这一点对应的实现位于 tilelang/tools/cuda/iket/init.py该模块将前端事件 APIfrontend、会话生命周期session、宿主侧输出辅助cli统一导出公开符号包括mark、range、range_push/range_pop、payload、session、enable/disable、profile_command、trace_files等。IKET 的接入完全走 TileLang 常规的targetcuda后端不依赖 TileScale 或 CuTe DSL 前端它利用 TileLang 的 CUDA 源码后处理回调tilelang_callback_cuda_postproc在生成 CUDA 文本后注入仪器化代码回调注册逻辑见 tilelang/tools/cuda/iket/session.py 的enable()函数回调实现入口为 tilelang/tools/cuda/iket/codegen.py 的inject_iket_cuda(...)。环境要求目标环境需要满足带 CUDA 支持的 TileLangCUDA 可用的 GPU 及驱动运行仓库自带示例时还需要 PyTorch外部的 IKET Python 包与运行时用于采集轨迹。可以用一条命令验证当前 Python 环境python -c import tilelang, iket, torchSM 架构适配TileLang 的代码生成路径同时支持 Hopper 之前与 Hopper 及以后的目标。从源码看codegen.py 中的_cluster_rank_instruction(...)会根据 target 解析出的 SM 版本决定 rank 指令SM90 及更新架构生成的仪器化 PTX 内联汇编读取%cluster_ctarankmov.b32 r, %cluster_ctarank;Hopper 之前的目标或无法判断目标架构时使用 cluster rank0mov.u32 r, 0;不会发射 Hopper 专属的寄存器。这一细节决定了 IKET 在 A100 与 H100 等平台上都能生成合法的汇编但集群维度行为不同。2. 快速上手在 iket.session 中构建并编译最小可用流程是在iket.session(...)上下文内构建并编译带仪器化的 kernelimport tilelang import tilelang.language as T from tilelang.tools.cuda import iket def instrumented_add(n: int, threads: int 128): T.prim_func def main( A: T.Tensor((n,), T.float32), B: T.Tensor((n,), T.float32), C: T.Tensor((n,), T.float32), ): with T.Kernel(T.ceildiv(n, threads), threadsthreads) as bx: with iket.range(block_total): for tx in T.Parallel(threads): i bx * threads tx if i n: iket.mark(before_store) C[i] A[i] B[i] iket.mark(after_store) return main with iket.session(output_dir/tmp/tilelang_iket): program instrumented_add(1024) kernel tilelang.compile( program, out_idx-1, targetcuda, execution_backendcython, )关键约束与细节会话必须在tilelang.compile(...)生成 CUDA 源码期间处于激活状态否则回调不会被触发生成的 CUDA 源码里不会有 IKET 元数据。推荐在会话内构建 kernel因为此时会话会拿到一份干净的事件注册表reset_eventsTrue会调用前端reset()清空事件 ID 分配但这不是必需的——事件元数据以 token 形式内嵌在 TIR 中因此进入会话之前构建好的PrimFunc在编译时仍携带全部所需信息原理见第 6 节。直接运行程序如仓库示例 examples/iket/minimal.py、examples/iket/all_features.py会编译并执行仪器化 kernel同时通过kernel.get_kernel_source()校验源码中是否包含__iket_meta_info与TL_IKET_EVENT符号并写出.cu源文件。要真正采集轨迹必须通过第 5 节的外部 IKET profiler 运行。仓库提供了三个由浅入深的可运行示例其覆盖范围在 examples/iket/README.md 中有说明示例文件覆盖内容examples/iket/minimal.py瞬时标记marker与一个 warp 局部 rangeexamples/iket/payload_minimal.py一个int32运行时 payload 标记examples/iket/all_features.pyrange、无 payload 标记、运行时 payload 标记的综合演示直接运行仅验证编译与执行不采集轨迹python examples/iket/minimal.py \ --iket-output-dir /tmp/tilelang_iket_minimal python examples/iket/payload_minimal.py \ --iket-output-dir /tmp/tilelang_iket_payload_minimal \ --iket-runtime-payloads python examples/iket/all_features.py \ --iket-output-dir /tmp/tilelang_iket_all_features \ --iket-runtime-payloads--iket-runtime-payloads是一个布尔开关示例内部会将其传给iket.session(runtime_payloads...)控制元数据是否以非 NoPayload形式发射见第 4 节。3. 事件 APIMarkers 与 Ranges3.1 三种基本用法瞬时事件用iket.mark(...)iket.mark(load_inputs)作用域用iket.range(...)作为 Python 上下文管理器包裹一段词法区域with iket.range(compute): # TileLang statements ...显式 push/pop API同样可用iket.range_push(compute) # TileLang statements iket.range_pop(compute)其中iket.range_start(...)与iket.range_end(...)分别是range_push(...)与range_pop(...)的别名见 frontend.py 第 109–116 行。range start 可以携带 payload但range-end 事件不能。3.2 warp 粒度与命名约束warp 粒度记录IKET 以 warp 为粒度记录 range。一个包含 4 个 warp 的 block对一个词法iket.range(...)作用域会产出4 条轨迹 range。名称长度事件与 range 名称上限为32 个 UTF-8 字节。源码中该限制由 metadata.py 的MAX_EVENT_NAME_BYTES 32定义前端_get_event(...)在注册时即校验超出直接抛ValueError。payload 类型冲突拒绝在同一个前端注册表中用不同的 payload dtype 复用同一 marker/range 名称会被拒绝。实现上frontend.py 的_validate_payload_compat(...)会比较首次注册的payload_type与当前使用的类型不一致即抛错。range 的 ID 派生range_id 取名称的 CRC32zlib.crc32因此同名的 range 跨 kernel 也具有一致的 range 标识而event_id由前端注册表顺序分配并且实现刻意跳过 31 这个 ID_RANGE_END_EVENT_ID 31保留给 range-end 专用事件见 frontend.py 第 16 行。push/pop 配对检查range_pop(...)使用线程局部栈校验配对pop 空栈或名称不匹配都会抛RuntimeError这有助于尽早发现不匹配的作用域书写错误。4. 运行时 Payload 机制4.1 支持的类型与描述符marker 和 range start 可以捕获一个 32 位标量值。TileLang 目前支持的 payload dtype 只有三种int32uint32float32对 TileLang 表达式应使用显式 payload 描述符iket.mark(store_index, payloadiket.payload(i, dtypeint32)) iket.mark(scale, payloadiket.payload(value, dtypefloat32))iket.payload(...)返回一个PayloadSpec表达式 dtype IKET 侧类型 ID的冻结 dataclass。简单 Python 标量int/float/bool和带dtype属性的表达式可以直接传入dtype 会被自动推断bool → uint32、int → int32、float → float32但显式指定 dtype 能让 trace 的 schema 无歧义。dtype 别名int/uint/float会被归一化到int32/uint32/float32其它类型直接触发NotImplementedError。4.2 运行时捕获是 opt-inwith iket.session(runtime_payloadsTrue): program instrumented_add(1024) kernel tilelang.compile(program, targetcuda)不开启runtime_payloadsTrue时的行为差异由 codegen.py 的_metadata_decls(...)与_event_macros(...)决定未开启payload schema 仍然编码在 TIR 元数据 token 里但发射的 IKET 元数据声明NoPayload生成的 payload 宏TL_IKET_EVENT_PAYLOAD_U32/F32退化为普通TL_IKET_EVENT(ID)即不写 payload 值。这样普通 marker 记录保持在 4 字节。已开启payload 事件会发射两条独立的 32 位 shared-memory store——一条写时间戳|事件 ID另一条写 payload 值详见第 6.3 节的宏实现。payload 值通过 IKET 的warp 级 dump 机制观察一个 payload 值通常代表该机制选中的 lane而不是 warp 内每个线程都上报同一个值。5. 采集与查看轨迹5.1 用外部 profiler 采集轨迹外部 profiler 通过 IKET 包的 CLI 暴露python -m iket.cli.main对综合示例 examples/iket/all_features.py 的完整采集命令rm -rf /tmp/tilelang_iket_all_features_profile python -m iket.cli.main \ --output-dir /tmp/tilelang_iket_all_features_profile \ --clobber \ profile \ --postprocess all \ -- \ python examples/iket/all_features.py \ --iket-output-dir /tmp/tilelang_iket_all_features_profile \ --iket-runtime-payloadsprofiler 会配置外部 IKET 运行时并在--之后启动目标命令。输出目录中可能包含iket_pid_0x....pftrace iket_pid_0x....pftrace.gz iket_pid_0x....trace.json iket_pid_0x....htmlTileLang 还可以从 Python 侧构造同一条 shell 命令command iket.profile_command( [python, examples/iket/all_features.py, --iket-runtime-payloads], directory/tmp/tilelang_iket_all_features_profile, ) print(command)profile_command(...)只返回带引号转义的命令字符串实现见 cli.py内部按python -m iket.cli.main --output-dir dir --clobber profile --postprocess all -- command拼装并用shlex.quote转义它本身不会启动 profiler。5.2 在 Perfetto 中查看生成的 HTML 需要能加载同目录下的 trace 文件因此应把输出目录作为静态站点服务cd /tmp/tilelang_iket_all_features_profile python3 -m http.server 8080然后在浏览器打开具体生成的文件例如http://localhost:8080/iket_pid_0x....html远程主机上用 SSH 端口转发ssh -L 8080:localhost:8080 userremote-host如果页面只显示 Perfetto 的落地页而没有数据请在 Perfetto UI 中手动导入对应的.pftrace文件。5.3 编程式解析 JSON 导出JSON 轨迹可以程序化检查例如提取所有store_indexmarker 的 payload 值import json from pathlib import Path trace_path max( Path(/tmp/tilelang_iket_all_features_profile).glob(*.trace.json), keylambda path: path.stat().st_size, ) data json.loads(trace_path.read_text()) launch data[launches][0] names data[stringTable] store_indices [ marker[payloadVal] for marker in launch[markers] if names[marker[markerNameIdx]] store_index and payloadVal in marker ] print(store_indices[:8])examples/iket/minimal.py 里也展示了如何用iket.trace_files(...)找到最新轨迹并统计launches/markers/ranges数量。6. 实现原理仪器化如何穿越编译6.1 会话参数与状态回滚iket.session(...)的完整签名with iket.session( reset_eventsTrue, overrideTrue, disable_on_exitTrue, output_dirNone, runtime_payloadsNone, disable_cacheTrue, ): ...各参数控制的宿主侧状态如下参数默认作用reset_eventsTrue清空此后构建 kernel 的前端事件分配已内嵌在PrimFunc里的元数据不受影响overrideTrue允许 IKET 在最外层会话激活期间替换已存在的tilelang_callback_cuda_postproc回调disable_on_exitTrue退出时恢复之前的回调限定作用域使用请保持默认output_dirNone创建目录、设置TL_IKET_OUTPUT_DIR环境变量并在会话期间配置 TileLang 的 IKET 路径辅助函数runtime_payloadsNone临时选择是否发射 payload 值None保持原有设置disable_cacheTrue默认绕过 TileLang 的KernelCache会话退出时恢复其先前状态从 session.py 的实现看_Session.__enter__会先对输出目录、payload 模式、缓存状态做快照再依次应用新状态并注册回调任何一步失败都会走except分支回滚已生效的部分。__exit__则按disable_on_exit恢复回调并回滚快照——这正是优先使用iket.session(...)的原因即使编译抛异常清理也一定发生。为什么disable_cache默认开启因为事件名与 schema 虽然是 TIR 缓存标识的一部分结构相同但事件名不同的 kernel 会得到不同的缓存 key但回调激活与 payload 模式是宿主侧编译状态。复用一份未带回调、或以不同 payload 模式编译的二进制会产生缺失或过期的仪器化。只有在调用方能控制这些条件时才应设disable_cacheFalse。回调是引用计数的enable()/disable()维护_enable_depth嵌套 IKET 会话保持外层回调激活退出最外层会话时才恢复 IKET 之前注册的回调输出目录、payload 模式、缓存状态同样在会话退出后恢复。更底层的手动生命周期 API 如下用于高级场景但需自行保证清理iket.enable() iket.is_enabled() iket.disable() iket.enable_runtime_payloads() iket.runtime_payloads_enabled() iket.disable_runtime_payloads()6.2 元数据 token事件信息内嵌于 TIR前端每次事件调用都会生成一个规范化的元数据 token并作为参数写入 TIR 中的call_extern调用。token 的构造见 metadata.py前缀固定为__tl_iket_v1_内容是{version, name, event_id, kind, range_id, payload_type, payload_iket_id}的 JSON经sort_keys压缩后做 URL-safe base64 编码解码端decode_event(...)对前缀、字段集合、版本、名称长度、kind 取值仅mark/range、ID 合法性event_id不能为 0 或 31逐项校验防止陈旧或损坏的元数据进入代码生成。这带来两个重要后果预构建的PrimFunc在会话进入与前端注册表重置之后仍保留事件元数据——事件信息不依赖进程级注册表结构相似但事件名不同的 kernel 具有不同的 IR 缓存标识——IKET 事件参与 TileLangKernelCache的 key 计算因此改事件名会触发重新编译而不是命中旧二进制。6.3 CUDA 源码后处理元数据数组与事件宏会话激活期间tilelang_callback_cuda_postproc触发 codegen.py 的inject_iket_cuda(...)完成四件事恢复事件用正则_EVENT_CALL_PATTERN从生成的 CUDA 中找出TL_IKET_EVENT_PAYLOAD_U32|_F32调用并解码 tokentoken 数量与事件调用数量必须一致否则报IKET metadata token is not attached to a supported event call。分配模块级事件 ID按(kind, name)归并重新分配从 1 开始的模块级 ID同样跳过 31冲突的元数据会直接抛错。发射 IKET 元数据数组每个事件生成一个 60 字节__device__常量数组事件 ID、instrument method、payload 类型 ID、range 信息、名称等range 另生成 72 字节数组含 CRC32 range_id外加一个 48 字节的__iket_meta_info头含 IKET magic 字节157, 241, 190, 186、条目数、事件上限等。这些数组以extern C声明、__attribute__((used, aligned(1)))保证不被优化掉是外部 IKET patcher 定位与解析事件的依据。定义 NativeDump 事件宏核心宏TL_IKET_EVENT(ID, ...)展开为一段 PTX 内联汇编.reg .b32 r, t; mov.b32 r, %cluster_ctarank; // 或 mov.u32 r, 0; mov.u32 t, %globaltimer_lo; or.b32 t, t, ID; mad.lo.u32 r, r, 0x1000000, 0x20; st.weak.shared.u32 [r], t; pmevent.mask ID;即以 cluster rank 计算 shared memory 槽位每 rank 预留0x1000000字节空间、偏移0x20起把globaltimer 低 32 位 | 事件 ID写入该槽位再触发pmevent.mask性能监控事件。无 payload 事件只写这一条 32 位记录。开启runtime_payloads后payload 宏改为两段独立汇编第一段写时间戳记录偏移0x20第二段把 payload 值以st.volatile.shared.u32写到偏移0x24最后触发pmevent.mask。payload store 特意使用volatile目的是阻止 ptxas 把这对 store 合并为STS.64——该 64 位记录形态是当前外部 IKET patcher不接受的。这是理解第 8 节故障排查的关键。6.4 宿主侧辅助函数CUDA 工具还提供了小型宿主辅助 API实现见 cli.pyiket.set_output_dir(/tmp/tilelang_iket) # 创建目录并设置 TL_IKET_OUTPUT_DIR iket.output_dir() # 返回当前输出目录 iket.output_path(kernel.cu) # 输出目录下的具体文件路径 iket.trace_files() # 返回 *.trace.json按大小降序 iket.profile_command([...], directory/tmp/tilelang_iket) # 只构造命令字符串这些辅助只管理路径与命令构造本身不采集轨迹。set_output_dir(...)会同时写进程变量与环境变量TL_IKET_OUTPUT_DIR示例脚本的--iket-output-dir默认值也正是回落到该环境变量。iket.event_table()返回最近构建 kernel 期间注册的事件列表name、event_id、kind、range_id、payload_type 等适合检查注册结果但不是代码生成的权威来源——权威来源是 TIR 中的 token。若需要显式重置前端事件分配调用iket.reset()。7. 限制Limitations文档明确列出当前 IKET 集成的边界使用与评估时应以此为准仅支持 TileLang 的CUDA 后端运行时 payload 仅限int32、uint32、float32事件与 range 名称上限32 UTF-8 字节不生成源码位置表trace 中的locIdx是 IKET 运行时的位置索引不是Python 或 TIR 行号IKET 以warp 粒度记录事件回调、payload 模式、事件注册表、kernel-cache 开关都是进程全局状态并行的编译工作流需要自行协调访问仓库测试 testing/python/tools/test_tilelang_tools_cuda_iket.py 的autousefixture 就专门做这种状态隔离与恢复;该集成依赖外部 IKET 运行时的私有元数据与 NativeDump 约定应视为实验性能力。8. 故障排查8.1 payload schema 存在但 trace 里没有payloadValschema 已内嵌但值未发射通常是没有在正确的 payload 模式下编译with iket.session(runtime_payloadsTrue): kernel tilelang.compile(program, targetcuda)同时检查该 marker 是否带有受支持的 payload 描述符iket.payload(...)且 dtype 属于 int32/uint32/float32。注意disable_cache默认开启就是为了避免命中未开 payload 模式时编译的旧二进制。8.2 profiler 在给 payload kernel 打补丁时失败payload 仪器化依赖两条独立的 32 位 store。可用反汇编检查生成二进制的记录形态nvdisasm kernel.cubin | grep -E STS|PMTRIG期望的形态是STS [addr], timestamp_with_event_id STS [addr0x4], payload_value PMTRIG event_id如果出现STS.64的时间戳/payload 对说明仪器化序列不再匹配 IKET 的 NativeDump 打补丁约定——检查编译时是否处于runtime_payloads会话、事件宏是否被后续处理改动以及外部 IKET 运行时版本是否与 TileLang 侧约定的宏实现兼容。9. 小结与延伸阅读IKET 把kernel 内的事件观测点做成了 TileLang 编译流水线的一等公民前端在 TIR 中内嵌自描述的元数据 token编译期后处理恢复 token、生成设备侧元数据数组与 PTX 事件宏外部 IKET 运行时再完成采样与轨迹导出。对需要剖析 TileLang CUDA kernel 内部阶段耗时例如 flash attention 中 load/compute/store 各阶段、block 级循环进度、关键索引值的开发者这套机制提供了比整体 kernel 计时更细的视角同时通过会话/引用计数/状态快照保证了与宿主编译状态的隔离与恢复。完整指南文档docs/tools/iket.md工具实现tilelang/tools/cuda/iket/init.py、tilelang/tools/cuda/iket/frontend.py、tilelang/tools/cuda/iket/session.py、tilelang/tools/cuda/iket/codegen.py、tilelang/tools/cuda/iket/metadata.py、tilelang/tools/cuda/iket/cli.py可运行示例examples/iket/README.md、examples/iket/minimal.py、examples/iket/payload_minimal.py、examples/iket/all_features.py行为测试回调注册/恢复、缓存、payload 元数据、SM90 rank 指令等testing/python/tools/test_tilelang_tools_cuda_iket.py【免费下载链接】tilelangDomain-specific language designed to streamline the development of high-performance GPU/CPU/Accelerators kernels项目地址: https://gitcode.com/GitHub_Trending/ti/tilelang创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考