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 diagnose
print(diagnose(triton_out, torch_ref))
** ▶ Open the 2‑minute Colab demo** (Runtime → GPU)
| You try… | What happens |
|---|---|
torch.allclose(...) |
|
False . No index. No cause. |
|
tl.device_print |
|
| Thousands of unsynced lines. No tensor context. | |
TRITON_INTERPRET=1 |
|
Breaks on bf16 / tl.load(tl.load(...)) . Not your real kernel. |
| Tool | What you get |
|---|---|
compare |
|
| Max-error indices, hotspots, neighborhoods, bf16-aware tolerances | |
bisect_axes |
|
| Minimal failing sub-tensor | |
bisect_tiles |
|
Failing program_id range on the compiled kernel |
|
classify_drift |
|
localized_bug vs reduction_order vs dtype_cast vs … |
|
ProbeBank |
|
| Fusion-boundary stage diffs (not print spam) |
from triton_blackhole import compare, bisect_axes, classify_drift, format_report
from triton_blackhole.classify import format_classification
print(format_report(compare(triton_out, torch_ref)))
print(format_classification(classify_drift(triton_out, torch_ref)))
print(bisect_axes(triton_out, torch_ref).report())
pip install triton-blackhole
Optional (Linux / WSL2 / Colab with NVIDIA):
pip install triton-blackhole[triton]
From source:
pip install -e ".[dev]"
| Piece | Native Windows | WSL2 + NVIDIA | Linux + NVIDIA | Colab GPU |
|---|---|---|---|---|
| Debugger (compare / bisect / classify) | ✅ | ✅ | ✅ | ✅ |
| Live Triton kernels | ❌ | ✅ | ✅ | ✅ |
No NVIDIA laptop? Use the Colab badge above.
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())
python
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
triton-blackhole compare triton_out.pt torch_ref.pt --bisect --suggest
We never interpret the kernel. We bisect:
Output space(tensor axes)** Grid space**(program_id
)Stage space(named intermediates)
Same binary as production — bf16, tensor cores, indirect loads included.
python examples/demo_softmax_drift.py
python examples/demo_tile_bisect.py
python examples/demo_triton_add.py # needs CUDA + triton
source & further reading
github.com — original article