triton.language.device_print

triton.language.device_print(prefix, *args, hex=False, _semantic=None)

Print the values at runtime from the device. String formatting does not work for runtime values, so you should provide the values you want to print as arguments. The first value must be a string, all following values must be scalars or tensors.

Calling the Python builtin print is the same as calling this function, and the requirements for the arguments will match this function (not the normal requirements for print).

tl.device_print("pid", pid)
print("pid", pid)

On CUDA, printfs are streamed through a buffer of limited size (on one host, we measured the default as 6912 KiB, but this may not be consistent across GPUs and CUDA versions). If you notice some printfs are being dropped, you can increase the buffer size by calling

triton.runtime.driver.active.utils.set_printf_fifo_size(size_bytes)

CUDA may raise an error if you try to change this value after running a kernel that uses printfs. The value set here may only affect the current device (so if you have multiple GPUs, you'd need to call it multiple times).

参数:

prefix -- a prefix to print before the values. This is required to be a string literal.
args -- the values to print. They can be any tensor or scalar.
hex -- print all values as hex instead of decimal

Example

import os
import torch
import torch_npu

import triton
import triton.language as tl

os.environ["TRITON_DEVICE_PRINT"] = "1"


@triton.jit
def device_print_kernel(x_ptr, BLOCK_SIZE: tl.constexpr):
    offsets = tl.arange(0, BLOCK_SIZE)
    x = tl.load(x_ptr + offsets)
    # device_print runs at runtime; the first argument must be a string prefix.
    # Set TRITON_DEVICE_PRINT=1 to see the output.
    tl.device_print("x =", x)


def test_device_print():
    BLOCK_SIZE = 8
    x = torch.randn(BLOCK_SIZE, device="npu", dtype=torch.float32)
    device_print_kernel[(1, )](x, BLOCK_SIZE=BLOCK_SIZE)
    torch.npu.synchronize()


if __name__ == "__main__":
    test_device_print()
    print("test_device_print PASSED!")

Special Restrictions

DataType: Ascend A2/A3 does not support uint16/uint32/uint64/fp64, Ascend 950 does not support fp64 (hardware limitation).
prefix: the first argument must be a string prefix; omitting it causes a compilation error.
Set environment variable TRITON_DEVICE_PRINT=1 to enable.