KTRACE: in-kernel wave tracing

KTRACE: in-kernel wave tracing#

Experimental. The API, record format, and output schema may change in a minor release. See API stability.

KTRACE records named phases inside kernels and exports a timeline for Perfetto. It currently supports gfx942 (CDNA3) and gfx950 (CDNA4), with one traced device per process. Each event records one active lane’s timestamp and optional payload for its wave.

Quick start#

Add coarse ranges to the kernel:

from flydsl.expr import ktrace

@flyc.kernel
def gemm(...):
    with ktrace.range("prologue"):
        ...
    with ktrace.range("mainloop"):
        ...
    with ktrace.range("epilogue"):
        ...

Enable tracing around the host launch, then open the exported gemm_ktrace.pftrace.gz at ui.perfetto.dev:

with ktrace.capture("gemm", blocks=(0, 0, 0)) as trace:
    launch_gemm(...)

print(trace.summary())
path = trace.export()

The trace nests tracks as kernel → launch → CU (XCD/SE/SH/CU) → WG(x, y, z) → WaveN. Each wave track starts with a Wave Span slice from its first to its last event; range_push ranges nest inside it and marks appear as instants. Token ranges (range_start/range_stop) may overlap, so they go on ranges N child tracks. An event’s payload appears as its payload argument. A range still open at the wave’s last event ends there and is labelled (unclosed). The block coordinates and wave index come from each wave’s buffer header; no annotation is needed.

blocks=(x, y, z) selects one workgroup without changing the launched grid. Start with one workgroup and a few phases. Each wave can record 255 events per launch; a range consumes two. Read or clear the trace between groups of launches to avoid filling the buffer.

The capture enables tracing for JIT specializations compiled inside it and reads the buffer on exit. summary() and export() use the same snapshot and return None when it contains no events. Calling them inside the capture keeps the records read so far. The snapshot remains available after exit. A CompiledFunction compiled before the capture runs untraced; the capture logs a warning when it records nothing.

Tracing is disabled by default. Disabled annotations emit no trace instructions. The first traced and untraced calls compile separate specializations; later calls reuse them. Leaving a capture restores the previous compile hints. A graph retains the specialization used during its construction.

See the runnable example:

python examples/07-ktrace_wave_timeline.py --layers 3

Kernel API#

Function	Purpose
`range(name, payload=None)`	Context manager for a named phase
`mark(name, payload=None)`	Point event
`range_push(name, payload=None)`	Open a stack range
`range_pop()`	Close the innermost stack range
`range_start(name, payload=None)`	Open a range and return its `RangeToken`
`range_stop(token, payload=None)`	Close a token range
`inactive_token(name)`	Inactive token for a loop-carried range

Use with ktrace.range(...) when both endpoints fit in one source block. range_push and range_pop pair in LIFO order along the executed path. A token range pairs on its token. The lanes that reach a range_stop can hold different tokens, as after a loop they leave at different iterations: the stop closes the latest range any of them started, and records nothing if none did.

A wave numbers its events in a register. An event in a lane-divergent branch takes the wave’s next slot. After every branch or loop that records, the lanes of the wave agree on the count, so later events continue past every slot it used, whichever lanes record them. When both branches of an if record, as in if lane < 32: mark("a") with else: mark("b"), the traced kernel runs the branches as two ifs on the same condition, the then branch first, so a wave that runs both keeps both records. Under a wave-uniform condition, such as the warp role, the agreement compiles to nothing and the second if adds one uniform branch.

Use tokens when the range crosses loop iterations:

stage = ktrace.inactive_token("k_tile")
for i in fx.range(0, k_tiles, 1):
    ab_full = ab_consumer.wait_and_advance()
    ktrace.range_stop(stage)
    stage = ktrace.range_start("k_tile")
    ...
ktrace.range_stop(stage)

This records intervals between wait boundaries. Sample iterations or use a single outer range when the loop would exceed the per-wave event limit.

Payloads#

Payloads accept Python bool, int, and float; FlyDSL integers up to 64 bits; and Float16, BFloat16, Float32, and Float64. Python integers must fit in signed or unsigned 64 bits. Python floats use 64 bits; half-precision values are promoted to float32.

Signedness and numeric values are preserved in the output. Non-finite floats are shown as nan, inf, or -inf. An omitted payload adds no argument; an explicit 0 remains 0. Use wave-uniform payloads: under divergence, the first active lane supplies the recorded value.

Captures and repeated launches#

A capture accumulates all launches until a read or clear. Repeated calls to the same kernel appear on separate timeline tracks.

Capture method	Effect
`clear()`	Synchronize and discard device records and the saved snapshot
`read()`	Synchronize, save a fresh snapshot, and reset the device buffer
`summary()`	Return per-phase totals for the snapshot
`launches()`	Return separate totals and metadata for each recorded launch
`export()`	Write a combined gzipped Perfetto trace
`export_launches(names=None)`	Write one gzipped Perfetto trace per recorded launch

with ktrace.capture("model", blocks=(0, 0, 0), out_dir="traces") as trace:
    for layer in layers:
        launch_layer(layer)

for launch in trace.launches():
    print(launch["launch_id"], launch["phases"])

paths = trace.export_launches(names=[f"layer_{i}" for i in range(len(layers))])

names must contain one string per recorded launch, in first-timestamp order. For overlapping streams, inspect the launch metadata before assigning layer names. ROCm includes HSA queue and dispatch IDs. Launches with no recorded events have no per-launch output. Capture windows cannot be nested.

For tracing enabled outside a capture, ktrace.summary() and ktrace.collect("kernel", out_dir=None) read and reset the device buffer. They share a cached snapshot until another traced launch or graph replay occurs. Untraced calls leave that snapshot available. Inside a capture they raise; use the capture’s methods.

Graph replay#

Enter the KTRACE context outside torch.cuda.graph(). Prepare the traced specialization before graph construction and pass an explicit stream through the launcher to every kernel launch:

import torch
import flydsl.expr as fx

with ktrace.capture("gemm", blocks=(0, 0, 0)) as trace:
    stream = torch.cuda.Stream()
    stream.wait_stream(torch.cuda.current_stream())
    with torch.cuda.stream(stream):
        launch_gemm(..., stream=fx.Stream(stream))

    graph = torch.cuda.CUDAGraph()
    with torch.cuda.stream(stream):
        with torch.cuda.graph(graph, stream=stream):
            launch_gemm(..., stream=fx.Stream(stream))

    trace.clear()
    graph.replay()

print(trace.summary())
print(trace.export())

An implicit default-stream launch does not enter a graph captured on PyTorch’s side stream. Entering, reading, or clearing a trace inside graph construction is invalid because these operations synchronize the device.

The graph retains its traced binary and buffer address after the context exits. Call trace.read() after a later replay to replace the saved snapshot. Read or clear between replay groups to stay within capacity. For multiple layers, capture each layer launch as a graph node; export_launches() also splits repeated calls to the same kernel. Run the example with --layers 3 --graph to try this.

Settings#

Variable	Default	Effect
`FLYDSL_KTRACE_ENABLE`	`0`	Enable instrumentation
`FLYDSL_KTRACE_BUFFER_BYTES`	`64 MiB`	Device buffer size; at most 2 GiB
`FLYDSL_KTRACE_BLOCKS`	all	`x,y,z` selects one workgroup; `xcc:N` selects one XCD
`FLYDSL_KTRACE_DUMP_DIR`	current directory	Trace output directory

An explicit ktrace compile hint overrides the environment. To set a persistent default for a launcher:

flyc.compile[{"ktrace": True}](launch, *args)

Use ktrace.capture(...) for a temporary setting. The buffer is allocated once per process. Keep its size consistent across compilation and execution; cached and deferred calls validate the compiled bound against the live allocation.

If independently cached kernels have conflicting event IDs, the JIT rebuilds the conflicting specialization and updates its cache. FLYDSL_RUNTIME_RUN_ONLY=1 prohibits that rebuild: prepare those traced kernels together once with run-only mode disabled.

Limits and interpreting timings#

255 events per wave per launch. Larger buffers and filtering other blocks do not increase this limit. Overflow raises during collection because dropped range endpoints would make the timeline incomplete.
Finite collection window. Waves reserve 256 slots each, including a header slot. A wave takes its slots from one of up to 256 partitions. Consecutive waves of a row of workgroups take consecutive partitions, so waves that start together do not contend for one counter and a launch fills the partitions evenly; successive launches start at different partitions. The default 64 MiB buffer has 64 partitions of about 128 waves, 8190 waves in all, and a trace overflows when its busiest partition fills: one launch fits up to about 8100 waves. Use block filters and periodic reads for large grids or repeated launches. The overflow error names a buffer size that holds the recorded launches.
One device per process. Access from another device raises. Trace multiple devices in separate processes. Traced binaries cannot be exported to C.
Synchronized collection. Reading waits for device work and serializes with traced host dispatches. Coordinate graph replays with reads and clears.
Reference-clock timestamps. s_memrealtime runs at 100 MHz (10 ns per tick), independently of the GPU core clock. Observed resolution on gfx950 is about 40 ns. Shorter ranges can appear as zero-duration markers.
Instrumentation changes timing. A wave makes one atomic when it starts, on its partition’s counter. An event reads the clock and issues two scalar stores, with no branch, no atomic and no VGPR. On MI355X an event in a loop cost about 20 ns, and ranges around the body, the K loop and the epilogue of a 4096-wave FP16 GEMM cost under 1%; a token range in every K-loop iteration cost about 2% and 16 more VGPRs. Timestamp consumption can drain outstanding LDS/SMEM traffic, and scheduling can move independent work across a timestamp. Use coarse phases and compare with an untraced baseline; subtracting a constant per-event cost does not undo these scheduling changes. Reading a trace idles the GPU while the host decodes it, so time traced launches between reads rather than across them.
Partial launch coverage. The first and last annotations need not coincide with kernel entry and exit. Gaps between tracks do not measure host launch latency or complete kernel overlap. Use a runtime profiler for those questions.