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/tilelang

IKET（IKET 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，该模块将前端事件 API（frontend）、会话生命周期（session）、宿主侧输出辅助（cli）统一导出，公开符号包括mark、range、range_push/range_pop、payload、session、enable/disable、profile_command、trace_files等。

IKET 的接入完全走 TileLang 常规的target="cuda"后端，不依赖 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 支持的 TileLang；
• CUDA 可用的 GPU 及驱动；
• 运行仓库自带示例时还需要 PyTorch；
• 外部的 IKET Python 包与运行时（用于采集轨迹）。

可以用一条命令验证当前 Python 环境：

python -c "import tilelang, iket, torch"

SM 架构适配

TileLang 的代码生成路径同时支持 Hopper 之前与 Hopper 及以后的目标。从源码看，codegen.py 中的_cluster_rank_instruction(...)会根据 target 解析出的 SM 版本决定 rank 指令：

• SM90 及更新架构：生成的仪器化 PTX 内联汇编读取%cluster_ctarank（mov.b32 r, %cluster_ctarank;）；
• Hopper 之前的目标，或无法判断目标架构时：使用 cluster rank0（mov.u32 r, 0;），不会发射 Hopper 专属的寄存器。

这一细节决定了 IKET 在 A100 与 H100 等平台上都能生成合法的汇编，但集群维度行为不同。

2. 快速上手：在 iket.session 中构建并编译

最小可用流程是：在iket.session(...)上下文内构建并编译带仪器化的 kernel：

import 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), threads=threads) 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, target="cuda", execution_backend="cython", )

关键约束与细节：

• 会话必须在tilelang.compile(...)生成 CUDA 源码期间处于激活状态，否则回调不会被触发，生成的 CUDA 源码里不会有 IKET 元数据。
• 推荐在会话内构建 kernel，因为此时会话会拿到一份干净的事件注册表（reset_events=True会调用前端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 中有说明：

直接运行（仅验证编译与执行，不采集轨迹）：

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. 事件 API：Markers 与 Ranges

3.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 取名称的 CRC32（zlib.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 只有三种：

• int32
• uint32
• float32

对 TileLang 表达式应使用显式 payload 描述符：

iket.mark("store_index", payload=iket.payload(i, dtype="int32")) iket.mark("scale", payload=iket.payload(value, dtype="float32"))

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-in

with iket.session(runtime_payloads=True): program = instrumented_add(1024) kernel = tilelang.compile(program, target="cuda")

不开启runtime_payloads=True时的行为差异（由 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-payloads

profiler 会配置外部 IKET 运行时，并在--之后启动目标命令。输出目录中可能包含：

iket_pid_0x....pftrace iket_pid_0x....pftrace.gz iket_pid_0x....trace.json iket_pid_0x....html

TileLang 还可以从 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 user@remote-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"), key=lambda 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_events=True, override=True, disable_on_exit=True, output_dir=None, runtime_payloads=None, disable_cache=True, ): ...

各参数控制的宿主侧状态如下：

从 session.py 的实现看，_Session.__enter__会先对输出目录、payload 模式、缓存状态做快照，再依次应用新状态并注册回调；任何一步失败都会走except分支回滚已生效的部分。__exit__则按disable_on_exit恢复回调并回滚快照——这正是"优先使用iket.session(...)"的原因：即使编译抛异常，清理也一定发生。

为什么disable_cache默认开启？因为事件名与 schema 虽然是 TIR 缓存标识的一部分（结构相同但事件名不同的 kernel 会得到不同的缓存 key），但回调激活与 payload 模式是宿主侧编译状态。复用一份"未带回调、或以不同 payload 模式编译"的二进制，会产生缺失或过期的仪器化。只有在调用方能控制这些条件时才应设disable_cache=False。

回调是引用计数的：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）逐项校验，防止陈旧或损坏的元数据进入代码生成。

这带来两个重要后果：

1. 预构建的PrimFunc在会话进入与前端注册表重置之后仍保留事件元数据——事件信息不依赖进程级注册表；
2. 结构相似但事件名不同的 kernel 具有不同的 IR 缓存标识——IKET 事件参与 TileLangKernelCache的 key 计算，因此改事件名会触发重新编译而不是命中旧二进制。

6.3 CUDA 源码后处理：元数据数组与事件宏

会话激活期间，tilelang_callback_cuda_postproc触发 codegen.py 的inject_iket_cuda(...)，完成四件事：

1. 恢复事件：用正则_EVENT_CALL_PATTERN从生成的 CUDA 中找出TL_IKET_EVENT_PAYLOAD_U32|_F32调用并解码 token；token 数量与事件调用数量必须一致，否则报IKET metadata token is not attached to a supported event call。
2. 分配模块级事件 ID：按(kind, name)归并，重新分配从 1 开始的模块级 ID（同样跳过 31），冲突的元数据会直接抛错。
3. 发射 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 定位与解析事件的依据。
4. 定义 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.py）：

iket.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 里没有payloadVal

schema 已内嵌但值未发射，通常是没有在正确的 payload 模式下编译：

with iket.session(runtime_payloads=True): kernel = tilelang.compile(program, target="cuda")

同时检查该 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 [addr+0x4], 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