如何使用triton.language.device_print打印数字？

我之前也踩过这个坑！在Triton 3.1.0版本里，tl.device_print确实有点“挑剔”——它只认字符串类型的输入，直接传整数、浮点数都会触发断言错误，而且你没法用Python的str()来转数字，因为设备端的代码是跑在CUDA/TPU硬件上的，根本不认Python的内置函数，这就是为什么你会碰到NameError。

不过别担心，Triton语言本身给咱们准备了解决方案：用tl.string_format()函数，这是专门为设备端设计的字符串格式化工具，能把数字转成符合要求的字符串，再传给tl.device_print就行。

比如你原来的代码可以改成这样：

import triton
import triton.language as tl

@triton.jit
def kernel():
    pid = tl.program_id(0)
    # 用%d格式化整数，把pid转成字符串
    tl.device_print(tl.string_format("当前Program ID：%d", pid))
    
    # 再给你个浮点数的示例
    float_val = tl.full([], 3.1415, dtype=tl.float32)
    tl.device_print(tl.string_format("浮点数示例：%.4f", float_val))

kernel[(1,)]()

这里的tl.string_format用法和咱们平时用Python的格式化逻辑很像：

• %d 对应整数类型
• %f 对应浮点数，还能加精度控制比如%.4f表示保留4位小数
• 第二个参数就是你要打印的数字变量或常量

运行这个修改后的代码，就能正常输出数字内容了。

另外提一句，如果之后你升级到更新版本的Triton，可能会发现tl.device_print已经支持直接打印数字了，但在3.1.0这个版本里，上面的方法是最靠谱的解决方案。

备注：内容来源于stack exchange，提问作者roz