## Technique library (KFORGE: techniques verified by the protected evaluator on training problems; use what fits this problem, ignore the rest)

### Verified technique: fp32-accurate GEMM on bf16 tensor cores (3-product split)
Literature: Henry, Tang & Heinecke 2019 (bf16x3 split products); Ootomo & Yokota 2022 (tensor cores with error correction beat the fp32 SIMT peak on A100). Checked here by the protected evaluator; facts, not orders.
- Physics of this A100 (strict fp32, measured on 3 GEMM shapes, x ~ U[0,1) and N(0,1)): eager fp32 cuBLAS (TF32 off) runs at 18-19 TFLOP/s. Plain TF32 tensor cores FAIL the tolerance (error 6-9x over it). Triton tl.dot(input_precision='tf32x3') passes but runs at only 21-22 TFLOP/s (1.12-1.20x). The 3-product bf16 split passes with max error <= 0.33 of the tolerance and runs at 64-79 TFLOP/s as ONE cuBLAS bf16 GEMM over a 3K-concatenated operand (3.5-4.2x), or 40-43 TFLOP/s as a Triton kernel that splits in registers.
- Verified on training problems (protected evaluator, correct, fast_1): L1_017 Matmul_with_transposed_B: 3.87x (cuBLAS split, 7.7 -> 2.0 ms); L2_076 Gemm_Add_ReLU: 3.55x (cuBLAS split, 8.0 -> 2.2 ms); L2_059 Matmul_Swish_Scaling: 2.21x (fused Triton split, 16.1 -> 7.3 ms); L2_076 Gemm_Add_ReLU: 2.23x (fused Triton split, 8.0 -> 3.6 ms); L3_001 MLP: 1.76x (fused Triton split, 10.7 -> 6.1 ms). Exact-fp32 kernels on such problems had reached ~1.0-1.3x.
- When it applies: the runtime is dominated by fp32 GEMM-like contractions (nn.Linear, matmul, bmm, einsum, attention QK^T and PV, 1x1 convolutions as GEMM) with K up to ~8192. Large M (>= ~1024): `linear_bf16x3` (cuBLAS over the concatenated operands); batched: the same concatenation with torch.bmm(a16, b16, out_dtype=torch.float32). Small M, or when an epilogue would cost an extra full pass: the fused Triton kernel. Keep the rest in fp32 and fuse epilogues (bias, activations, normalisation) into one pass over the GEMM output.
- Numerics: x = xh + xl with xh = bf16(x), xl = bf16(x - xh); x.w ~= xh.wh + xh.wl + xl.wh with fp32 accumulation (the dropped xl.wl term is ~2^-16 relative). Error grows ~sqrt(K); for K >> 8192 or cancellation-heavy outputs, CHECK and read max_abs_diff. bf16 keeps the fp32 exponent range (no scaling needed). Dropping the cross terms is plain bf16 and fails.
- Building block (verified; copy and adapt; `split3` is your own kernel launched every forward):
```python
# P1: fp32-accurate GEMM on bf16 tensor cores (3-product split, Henry/Tang/Heinecke 2019; Ootomo/Yokota 2022).
# x = xh + xl with xh = bf16(x), xl = bf16(x - xh) (same for w); x.w ~= xh.wh + xh.wl + xl.wh (drops xl.wl ~ 2^-16).
# Measured on A100: error <= 0.34 of atol=rtol=1e-4 at K <= 8192; cuBLAS fp32 = 19 TFLOP/s.
#   linear_bf16x3(x, w)         one cuBLAS bf16 GEMM over a 3K-concatenated operand, fp32 out: 65-79 TFLOP/s (3.5-4.2x)
#   linear_bf16x3_fused(x, w, b, ACT, s1, s2)  Triton, splits inside the kernel, fused bias+epilogue: 40-45 TFLOP/s
#       (2.2-2.6x); use it when M is small (W re-read dominates) or the epilogue would cost an extra pass.
import torch
import triton
import triton.language as tl


@triton.jit
def split3_kernel(src, dst, R, K, LO_MID: tl.constexpr, BR: tl.constexpr, BKK: tl.constexpr):
    # src [R, K] fp32 contiguous -> dst [R, 3K] bf16: [hi | hi | lo] (LO_MID=0, left operand) or [hi | lo | hi] (right)
    r = tl.program_id(0) * BR + tl.arange(0, BR)
    k = tl.program_id(1) * BKK + tl.arange(0, BKK)
    m = (r[:, None] < R) & (k[None, :] < K)
    v = tl.load(src + r[:, None] * K + k[None, :], mask=m, other=0.0)
    h = v.to(tl.bfloat16)
    lo = (v - h.to(tl.float32)).to(tl.bfloat16)
    base = dst + r[:, None] * (3 * K) + k[None, :]
    tl.store(base, h, mask=m)
    if LO_MID:
        tl.store(base + K, lo, mask=m)
        tl.store(base + 2 * K, h, mask=m)
    else:
        tl.store(base + K, h, mask=m)
        tl.store(base + 2 * K, lo, mask=m)


def split3(t, lo_mid):
    t = t.contiguous()
    R, K = t.shape
    out = torch.empty((R, 3 * K), device=t.device, dtype=torch.bfloat16)
    split3_kernel[(triton.cdiv(R, 32), triton.cdiv(K, 128))](t, out, R, K, LO_MID=lo_mid, BR=32, BKK=128)
    return out


def linear_bf16x3(x2d, w):
    """x2d [M, K] fp32, w [N, K] fp32 (nn.Linear layout) -> x2d @ w.T in fp32 (no bias). For A @ B with B [K, N],
    pass w = B.t() (split3 makes it contiguous)."""
    return torch.mm(split3(x2d, 0), split3(w, 1).t(), out_dtype=torch.float32)


@triton.autotune(
    configs=[
        triton.Config({"BM": 128, "BN": 128, "BK": 32, "GM": 8}, num_stages=3, num_warps=8),
        triton.Config({"BM": 128, "BN": 256, "BK": 32, "GM": 8}, num_stages=3, num_warps=8),
        triton.Config({"BM": 64, "BN": 128, "BK": 32, "GM": 8}, num_stages=4, num_warps=4),
        triton.Config({"BM": 64, "BN": 64, "BK": 64, "GM": 8}, num_stages=4, num_warps=4),
    ],
    key=["M", "N", "K"],
)
@triton.jit
def gemm_bf16x3_fused_kernel(x_ptr, w_ptr, b_ptr, y_ptr, M, N, K, sxm, sxk, swn, swk, sym, syn, s1, s2,
                             HAS_BIAS: tl.constexpr, ACT: tl.constexpr,
                             BM: tl.constexpr, BN: tl.constexpr, BK: tl.constexpr, GM: tl.constexpr):
    # y = act(x @ w.T + b); ACT: 0 none, 1 relu, 2 swish(v)*s1, 3 v*s1 + s2 (affine). Add your own epilogue here.
    pid = tl.program_id(0)
    npm = tl.cdiv(M, BM)
    npn = tl.cdiv(N, BN)
    width = GM * npn
    gid = pid // width
    fm = gid * GM
    gsz = min(npm - fm, GM)
    pm = fm + (pid % width) % gsz
    pn = (pid % width) // gsz
    rm = pm * BM + tl.arange(0, BM)
    rn = pn * BN + tl.arange(0, BN)
    rk = tl.arange(0, BK)
    xp = x_ptr + rm[:, None] * sxm + rk[None, :] * sxk
    wp = w_ptr + rn[None, :] * swn + rk[:, None] * swk
    acc = tl.zeros((BM, BN), dtype=tl.float32)
    acc2 = tl.zeros((BM, BN), dtype=tl.float32)
    for k in range(0, tl.cdiv(K, BK)):
        km = (k * BK + rk) < K
        a = tl.load(xp, mask=(rm[:, None] < M) & km[None, :], other=0.0)
        b = tl.load(wp, mask=km[:, None] & (rn[None, :] < N), other=0.0)
        ah = a.to(tl.bfloat16)
        al = (a - ah.to(tl.float32)).to(tl.bfloat16)
        bh = b.to(tl.bfloat16)
        bl = (b - bh.to(tl.float32)).to(tl.bfloat16)
        acc2 = tl.dot(ah, bl, acc2)
        acc2 = tl.dot(al, bh, acc2)
        acc = tl.dot(ah, bh, acc)
        xp += BK * sxk
        wp += BK * swk
    v = acc + acc2
    if HAS_BIAS:
        v = v + tl.load(b_ptr + rn, mask=rn < N, other=0.0)[None, :]
    if ACT == 1:
        v = tl.maximum(v, 0.0)
    elif ACT == 2:
        v = v * tl.sigmoid(v) * s1
    elif ACT == 3:
        v = v * s1 + s2
    tl.store(y_ptr + rm[:, None] * sym + rn[None, :] * syn, v, mask=(rm[:, None] < M) & (rn[None, :] < N))


def linear_bf16x3_fused(x2d, w, b=None, act=0, s1=1.0, s2=0.0):
    M, K = x2d.shape
    N = w.shape[0]
    y = torch.empty((M, N), device=x2d.device, dtype=torch.float32)
    grid = lambda META: (triton.cdiv(M, META["BM"]) * triton.cdiv(N, META["BN"]),)
    gemm_bf16x3_fused_kernel[grid](x2d, w, b if b is not None else y, y, M, N, K, x2d.stride(0), x2d.stride(1),
                                   w.stride(0), w.stride(1), y.stride(0), y.stride(1), float(s1), float(s2),
                                   HAS_BIAS=b is not None, ACT=act)
    return y
```


## Kernel memory (items verified by the protected evaluator on training problems; may not all apply)

### Lessons about this environment (from official checker verdicts and verified outcomes)
- Pitfall: A reply without a ```python fenced code block was rejected by the checker. Always wrap the full ModelNew source in a ```python block.
- Checker rule: getattr needs a literal, non-private attribute name
- Checker rule: forward modified its input tensors in place
- Checker rule: do not define a class named Model (the reference); define ModelNew only
- Process: never SUBMIT code that has not passed ACTION: CHECK in this episode. Untested submissions that fail to compile or mismatch the reference score 0; a CHECK costs one reply and shows the compiler output.
- Pitfall: cpp_sources is compiled by plain c++ without CUDA headers: never use cudaStream_t or at::cuda::getCurrentCUDAStream there. Put wrappers needing streams in cuda_sources, keeping only plain-type prototypes in cpp_sources.
- Pitfall: If cpp_sources must reference CUDA types or the current stream, explicitly #include <cuda_runtime.h> and <ATen/cuda/CUDAContext.h>; torch/extension.h does not provide them.
- Checker rule: assignment to attribute 'm.conv1' is not allowed (only self.<x>)
- Checker rule: assignment to attribute 'm.bn1' is not allowed (only self.<x>)
- Checker rule: assignment to attribute 'm.conv2' is not allowed (only self.<x>)
- Checker rule: assignment to attribute 'm.bn2' is not allowed (only self.<x>)
- Checker rule: assignment to attribute 'm.downsample' is not allowed (only self.<x>)
- Checker rule: At least one of your own kernels must actually launch during forward; a pure torch path (even with a compiled but unused kernel) is rejected as no_custom_kernel.
- Pitfall: A reply without a ```python fenced code block was rejected outright; always emit the full ModelNew inside a single ```python block.

### Pitfalls seen on earlier network problems (wrong results or slower than eager)
- Pitfall: Don'...'s fused path alone.
- Pitfall: In training mode, don't replace fused cuDNN BatchNorm with Python-level mean/var reductions plus separate scale/shift ops; each adds kernel launches and memory passes, which makes the result much slower than eager.
- Pitfall: Avoid extra .contiguous()/torch.empty_like allocations and per-call Triton launch overhead on small tensors; launch overhead often exceeds savings from fusing two cheap elementwise ops.
- Pitfall: In conv-dominated networks, fusing only a tiny elementwise op (add+ReLU) via Triton can't offset launch/Python overhead; fuse BatchNorm into conv weights (eval mode) or target the dominant ops instead.
- Pitfall: In training-mode BN, computing batch mean and var with separate torch reductions plus running-stat updates adds full memory passes. Eager's fused cudnn batch_norm is hard to beat this way; benchmark per layer before replacing.
- Pitfall: Don't fuse only BN+ReLU after cuDNN convs while computing BN scale/shift with many small eager ops each forward; the extra launches and passes outweigh the fusion. Fold BN into conv weights or fuse into the conv epilogue.