passnet-triton-opt

---

name: passnet-triton-opt description: > Make a single PassNet Triton kernel fast AND numerically faithful: performance model (when a replacement can win at all), block/grid/warp tuning, autotune policy, launch overhead, and per-op numeric recipes to pass the dtype baseline tolerances. Use when a pass matches and is correct but speedup below expected, or when correctness fails by small numeric margins.

1. Performance model — know if a win is even possible

Measured costs (A100, torch 2.7.1 / triton 3.3.1; same order of magnitude on the eval box):

component	cost per call
dynamo guards + FX graph interpretation + dispatch wrapper	~25–45 µs on a 3-op graph; grows with node count (~150–230 µs seen on a 13-node graph) — calibrate from your first eval
Triton kernel Python launch (kernel[grid](...))	~19 µs each
torch.empty_like / torch.empty	~2–3 µs each
each aten op you REMOVED from the graph	+5–25 µs back (small ops) or its real kernel time

So per variant: compiled_e2e ≈ eager_e2e − Σ(absorbed aten µs) + tax + Σ(your kernels µs). Rules of thumb:

• Whole eager forward < ~150 µs and you can absorb ≤2 cheap ops → ceiling < 1. Ship a safe correct pass, accept ~0.8–0.95, spend your time on another bottleneck.
• Small or narrow fusions can match and be numerically correct yet still lose e2e when launch/framework overhead exceeds the eager work removed.
• Layout-only, view/reshape-only, split/fanout, tiny activation, and small reduction-like kernels can also lose even when their local kernel is correct: eager may treat the work as metadata, use a vendor-tuned primitive, or pay less framework overhead than the compiled path. Treat these as floor/downside caps unless they bridge substantial compute or remove an actual materialization.
• The lever is ops absorbed per launch (and avoiding extra launches/allocs), not shaving 10% off an already-tiny kernel.
• Big tensors (≥ a few MB; kernel ≥ 100 µs): overhead noise vanishes; now memory traffic is everything — a clean fused kernel that reads each input once and writes once usually beats the multi-pass eager chain by the ratio of memory passes saved.
• The biggest wins can be ALGORITHMIC, not tuning — but the default for a heavy conv/matmul is still "leave it in the library, fuse its tail," because the library is near-optimal inside the regime it tuned for. The clearest case where a rewrite beats it: a non-overlapping convolution (stride == kernel_size, kernel_size > 1, padding == 0) is a dense patch matmul while cuDNN runs a heavier general im2col path — a tl.dot patch-gather kernel (kernel-templates §13) can win big. It usually does NOT pay to rewrite a 1×1 conv (already a cuBLAS-optimal pointwise GEMM), an overlapping conv (im2col blows up memory), a depthwise/grouped conv, or a general dense matmul — those are vendor sweet spots. Treat this as a prior, not a ban: the completed eval decides. Pursue any rewrite only after a floor is banked, and don't dismiss the non-overlapping case as "cuDNN, unbeatable" without checking stride == kernel_size.
• For large cat/stack/gather-style outputs, one giant kernel is not automatically best if it needs heavy branching or irregular index logic. A small fixed number of simple kernels can be faster and more timing-stable when each launch writes a large contiguous region. Compare sequential micro-bench signals, then trust the completed evaluation.
• Repeated round evidence favors long memory-bound elementwise and broadcast-affine chains when one legal single-output kernel absorbs multiple affine, activation, residual, layout, or normalization consumers. These are the first places to spend tuning effort after a correct floor because saved memory passes can amortize wrapper and launch overhead.
• Repeated completed evaluations also favor normalization-affine compute prefixes and concat-to-normalization-affine prefixes when they remove large memory passes in one launch. Treat kernel layout choices for concat branches as workload-dependent: a broad prefix can win strongly, but small/branchy cases or attempts to absorb extra producers can regress in score, correctness, or timing stability. Keep the best completed floor/prefix state and revert completed regressions instead of stacking more small kernels or sweeping blindly.
• Reduction-only regions and vendor-tail replacements still need completed-evaluation proof. A correct reduction kernel, a small tail behind a vendor op, or a region that replaces a vendor-tuned primitive can remain below eager even when local kernel timing looks plausible. Use the completed eval to decide keep/revert; do not infer success from the triage family.
• Compare [Speedup][gpu] vs [Speedup][e2e] in the eval log: gpu≈e2e&both<1 → on-stream gaps/launches dominate (absorb more ops / fewer launches); gpu>1 but e2e<1 → host-side overhead (reduce launches, drop autotune key churn, simplify wrapper Python).
• Only a completed evaluation can prove the final overhead-bound conclusion. Local micro-bench and check_pattern --bench are pre-flight signals; if the completed eval is matched and correct but below eager for the maximal legal region, stop or revert instead of sweeping launch-bound tuning knobs.

2. Kernel quality checklist (in order of impact)

1. One read per input element, one write per output element. No multi-launch chains when one kernel can hold the chain in registers. Exception: for heterogeneous direct-output cat/stack branches, compare a branch-heavy single kernel with grouped launches; choose by measured stable e2e speed, not by launch count alone.
2. Contiguous innermost access: thread offs should walk the last (stride-1) dim. For strided/permuted inputs pass strides; never .contiguous() (poison-illegal) — absorbing the transpose INTO the index math is free.
3. Block sizing: elementwise → BLOCK 1024 (n ≥ 1M: 2048, num_warps=8). Row kernels → BLOCK_N = next_power_of_2(N), and start with num_warps=4 — more warps on small rows HURTS (measured: 512-elem rows, warps=8 was 1.8× slower than warps=4). If N > 16384 switch to a looped accumulation (template 6 in kernel-templates). For tiny rows (≤1024 elems), SWEEP rows-per-program ∈ {1, 2, 4, 8} (2D tile) — the winner is shape-dependent (measured 1.45× for ROWS=4 at D=512/131k rows, but ROWS=1 best on a 768-wide LayerNorm); never assume, always micro-bench (§5).
4. Grid: prefer 1D cdiv(n, BLOCK); for row kernels one program per row is fine up to ~256k rows. Avoid grids of 1–4 programs on big inputs (underutilization) — split.
5. num_stages: only matters for tl.dot/looped kernels (3–4); elementwise ignore.
6. Avoid recompilation churn: every distinct tl.constexpr value (incl. BLOCK chosen from shape) compiles a new binary during warmup. 2–3 specializations fine; per-shape formulas over 100-variant samples → timeout risk. Derive constexprs from a small fixed set (e.g. BLOCK_N = next_pow2 clamped to {128, 512, 1024, 4096}).
7. Autotune policy: @triton.autotune is legal and runs during the untimed warmup (poison-safe). Use ≤4–6 configs, key=['N']-style coarse keys. SKIP autotune when the sample has many graph variants (tuning cost × variants × 600 s timeout) or when the kernel is launch-bound anyway. A hand-picked config from one local micro-benchmark is usually within 5%.
8. Scalar args beat tensor args: precompute n, C, HW, strides in the wrapper (metadata access is poison-legal) — never load shape info from tensors in-kernel.
9. Reduce wrapper Python: no dict lookups in hot path beyond the route if; compute grid inline; allocate with torch.empty (not zeros — don't pay a memset you'll overwrite).

3. Numeric fidelity — passing t=-5 baseline tolerances

Baselines: fp32 rtol 1.3e-6/atol 1e-5 (tight — mirror eager's algorithm), fp16 rtol 1e-3, bf16 rtol 1.6e-2 (loose). The same kernel must satisfy ALL dtype variants.

Recipes (these match aten's CUDA implementations):

• Always compute in fp32 (.to(tl.float32) after load), store with .to(out_ptr.dtype.element_ty). For fp16/bf16 inputs this reproduces aten (which also accumulates in fp32) to ≤1 ulp.
• Reductions: tl.sum/tl.max over a row in one block ≈ aten's order closely enough for fp32 rtol 1.3e-6 at typical sizes (N ≤ 16k). For looped accumulation keep ONE fp32 accumulator (don't reorder into multiple partial schemes unnecessarily).
• softmax: subtract row max (aten does), exp, divide by sum. Don't use exp2 shortcuts on fp32 variants.
• layer_norm/batch_norm: biased variance (÷N), 1/sqrt(var+eps) in fp32, then scale/shift.
• mean: sum/N (single division at the end, not running mean).
• GELU exact: 0.5x(1+erf(x/√2)) via tl.math.erf; tanh GELU only when model.py says approximate='tanh'.
• sigmoid/silu: 1/(1+exp(-x)) fp32 — fine at all baselines.
• x.pow(2) → x*x; torch.rsqrt → 1.0/tl.sqrt (matches aten rsqrt within 1 ulp; if a strict fp32 variant complains, try tl.math.rsqrt).
• Integer/bool tensors: compare with torch.equal (atol/rtol=0 conceptually) — your kernel must be EXACT (indices, masks, argmax-style ops: replicate tie-breaking by first-occurrence).
• Output dtype must equal eager's per output ([Datatype] check), including intermediate .to(torch.float32) nodes you absorbed — store in the dtype eager would have at the region's output.

Diagnosing accuracy failures from eval logs: [Correctness][max_diff] magnitude vs the table above (parse_eval_log.py does this) —

• max_diff ~1e-3·|values| on fp16 only → fp16 rounding mismatch: ensure fp32 internal math.
• fails ONLY fp32 variants → algorithmic mismatch (reduction order, formula) — mirror aten's exact sequence.
• max_diff huge (≥1) → indexing/masking bug, not numerics: check strides, other= values contaminating reductions (use -inf for max, 0 for sum, and mask BEFORE divide), uninitialized output regions (mask coverage), offs overflow on >2^31 elements (use tl.int64 offsets when numel > 2e9).
• NaN/Inf only in compiled → division by masked-out zero lanes; tl.where the mask before the division, not after.

4. Launch-overhead reduction (when e2e-bound)

• Merge kernels (one launch per region; the whole point of fusion).
• Do not add a tiny independent kernel after a correct-but-slow state unless it absorbs enough real eager work to pay for another launch. If a completed eval shows the added region lowers score, revert to the best completed state and stop that line.
• Avoid per-call next_power_of_2 recompute churn → precompute in wrapper; constexpr set small (see 2.6).
• Don't return tuples/lists from the wrapper unless the pattern output needs it.
• Buffer reuse across calls (cache out keyed by shape in a module dict) saves ~2 µs but risks aliasing bugs if the model returns your output AND mutates downstream — only as a last resort, verify with full eval.
• CUDA graphs / torch.compile of the wrapper / no_dispatch tricks: NOT allowed (torch.* calls blocked / hacking). Don't.

5. Micro-benchmark before full eval (GPU, seconds)

import torch, time
def bench(fn, *args, iters=500):
    for _ in range(20): fn(*args)
    torch.cuda.synchronize(); t0 = time.perf_counter()
    for _ in range(iters): fn(*args)
    torch.cuda.synchronize(); return (time.perf_counter() - t0) / iters * 1e6  # µs

Compare: (a) eager region ops chained, (b) your wrapper. Build inputs with the shapes / dtypes from weight_meta (use the checker's replay). Target: wrapper ≤ eager_region − 40 µs for small graphs; for ms-scale kernels target the memory-traffic bound (bytes moved / ~1.5 TB/s on A100 ≈ achievable µs).

check_pattern.py --bench (passnet-feedback) automates exactly this comparison per variant; trust the full eval for the final number (it adds guards/FX costs that local micro-bench misses).