triton-blackhole

torch.allclose failed. Now what?

triton-blackhole is a deterministic numerical debugger for Triton kernels. It finds where your output diverges from a PyTorch reference — and whether it's benign fp16/bf16 drift or a real bug — without tl.device_print floods or TRITON_INTERPRET.

from triton_blackhole import verify_drift

@verify_drift(torch_ref, block_sizes=(BLOCK_M, BLOCK_N))
def run(a, b):
    return launch_triton(a, b)

run(a, b)  # silent if OK; AssertionError + drift artifact if not

▶ Open the 2‑minute Colab demo (Runtime → GPU)

The pain

The fix (TritonDrift loop)

from triton_blackhole import verify_drift, run_drift_verify, index_to_program_id

# Decorator (pytest-friendly)
@verify_drift(torch_ref, block_sizes=(32, 32), raise_on_fail=True)
def run(a, b):
    return triton_launch(a, b)

# Or functional
art = run_drift_verify(tri_out, ref_out, block_sizes=(32, 32))
print(art.report())  # compare + grid map + bisect

Optional: AST dump on the failing block

def relaunch(ikernel, failing_pid, debug_buf, a, b):
    out = torch.empty_like(...)
    ikernel[(grid,)](
        a, b, out, ...,
        _bh_dbg_ptr=debug_buf,
        _BH_FAILING_PID=failing_pid,
    )
    return out

@verify_drift(
    torch_ref,
    block_sizes=(BLOCK_M, BLOCK_N),
    kernel=my_kernel,          # original @triton.jit fn
    probes=["acc"],            # local names assigned in the kernel
    relaunch=relaunch,
)
def run(a, b):
    return launch(a, b)

Install

pip install triton-blackhole

Optional (Linux / WSL2 / Colab with NVIDIA):

pip install triton-blackhole[triton]
# On Colab, pin Triton to whatever torch wants, e.g.:
# pip install "triton==3.6.0"

From source:

pip install -e ".[dev]"

Platforms

No NVIDIA laptop? Use the Colab badge above.

Tile bisection (real kernels)

from triton_blackhole import bisect_tiles

def launch(pid_lo, pid_hi):
    return run_triton_kernel(..., pid_lo=pid_lo, pid_hi=pid_hi)

print(bisect_tiles(launch, torch_ref, num_programs=grid).report())

Stage probes (not device_print)

from triton_blackhole.probe import ProbeBank

bank = ProbeBank()
bank.capture("pre_softmax", scores_ref, side="ref")
bank.capture("pre_softmax", scores_tri, side="tri")
print(bank.report())  # first diverging stage

CLI

triton-blackhole compare triton_out.pt torch_ref.pt --bisect --suggest

Why not TRITON_INTERPRET?

We never interpret the kernel. We bisect:

1. Output space (tensor axes)
2. Grid space (program_id)
3. Stage space (named intermediates)

Same binary as production — bf16, tensor cores, indirect loads included.

Examples

python examples/demo_verify_drift.py
python examples/demo_verify_ast_triton.py  # needs CUDA + triton
python examples/demo_softmax_drift.py
python examples/demo_tile_bisect.py
python examples/demo_triton_add.py

License

MIT · brian-mwirigi/triton-blackhole