## 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 matmul problems (wrong results or slower than eager)
- Pitfall: For memory-bound broadcast or elementwise ops, eager is already near peak bandwidth. Use 2D row/column tiling instead of per-element integer division on flat indices, and load each broadcast operand once per tile.
- Pitfall: When using split-K, keep the split count small; a very large split writes and re-reads a huge partial-sum buffer plus an extra reduction launch, eating the occupancy gain.
- Pitfall: For bandwidth-bound Triton kernels, give each program more work: larger tiles, multiple rows per program, more warps. Tiny per-program work and Python launch overhead lose to cuBLAS/ATen kernels.
- Pitfall: Don't expect a single-pass memory-bound elementwise or broadcast kernel to beat eager, which already runs near bandwidth peak. Win by fusing multiple ops or reducing memory traffic.
- Pitfall: If writing a Triton GEMM anyway, use @triton.autotune over several BM/BN/BK, num_warps and num_stages configs, plus grouped tile ordering. A single fixed 64x64x32 config is far slower than eager.
- Pitfall: Avoid heavy split-K with fp32 atomic_add plus torch.zeros init; it adds memory traffic and is nondeterministic. Prefer few splits with a separate reduction pass, and autotune tile sizes.

### A verified kernel for a related matmul problem (L1_015 Matmul_for_lower_triangular_matrices, 4.58x): adapt what applies
```python
import torch
import torch.nn as nn
import triton
import triton.language as tl

@triton.jit
def tril_mm_kernel(A_ptr, B_ptr, C_ptr, N,
                   BLOCK_M: tl.constexpr, BLOCK_N: tl.constexpr, BLOCK_K: tl.constexpr):
    pid_m = tl.program_id(0)
    pid_n = tl.program_id(1)

    m_lo = pid_m * BLOCK_M
    n_lo = pid_n * BLOCK_N

    if n_lo > (m_lo + BLOCK_M - 1):
        return

    offs_m = m_lo + tl.arange(0, BLOCK_M)
    offs_n = n_lo + tl.arange(0, BLOCK_N)
    offs_k = tl.arange(0, BLOCK_K)

    acc = tl.zeros((BLOCK_M, BLOCK_N), dtype=tl.float32)

    k_start = (n_lo // BLOCK_K) * BLOCK_K
    k_end = tl.minimum(m_lo + BLOCK_M, N)

    k = k_start
    while k < k_end:
        kk = k + offs_k
        a = tl.load(A_ptr + offs_m[:, None] * N + kk[None, :],
                    mask=(offs_m[:, None] < N) & (kk[None, :] < N), other=0.0)
        b = tl.load(B_ptr + kk[:, None] * N + offs_n[None, :],
                    mask=(kk[:, None] < N) & (offs_n[None, :] < N), other=0.0)
        acc += tl.dot(a, b, allow_tf32=False)
        k += BLOCK_K

    mask_tri = offs_m[:, None] >= offs_n[None, :]
    acc = tl.where(mask_tri, acc, 0.0)
    store_mask = (offs_m[:, None] < N) & (offs_n[None, :] < N)
    tl.store(C_ptr + offs_m[:, None] * N + offs_n[None, :], acc, mask=store_mask)


class ModelNew(nn.Module):
    def __init__(self):
        super().__init__()

    def forward(self, A, B):
        A = A.contiguous()
        B = B.contiguous()
        N = A.shape[0]
        C = torch.zeros_like(A)
        BLOCK_M = 64
        BLOCK_N = 64
        BLOCK_K = 32
        grid = (triton.cdiv(N, BLOCK_M), triton.cdiv(N, BLOCK_N))
        tril_mm_kernel[grid](A, B, C, N, BLOCK_M=BLOCK_M, BLOCK_N=BLOCK_N, BLOCK_K=BLOCK_K,
                             num_warps=4, num_stages=3)
        return C
```

### A verified kernel for a related matmul problem (L1_014 Matmul_for_upper_triangular_matrices, 4.11x): adapt what applies
```python
import torch
import torch.nn as nn
import triton
import triton.language as tl

@triton.jit
def triu_mm_kernel(A_ptr, B_ptr, C_ptr, N,
                   BLOCK_M: tl.constexpr, BLOCK_N: tl.constexpr, BLOCK_K: tl.constexpr):
    pid_m = tl.program_id(0)
    pid_n = tl.program_id(1)

    m_lo = pid_m * BLOCK_M
    n_lo = pid_n * BLOCK_N

    # skip tiles entirely in lower triangle: n_hi < m_lo
    if (n_lo + BLOCK_N - 1) < m_lo:
        return

    offs_m = m_lo + tl.arange(0, BLOCK_M)
    offs_n = n_lo + tl.arange(0, BLOCK_N)
    offs_k = tl.arange(0, BLOCK_K)

    acc = tl.zeros((BLOCK_M, BLOCK_N), dtype=tl.float32)

    # k from i to j => k in [m_lo, n_hi]
    k_start = (m_lo // BLOCK_K) * BLOCK_K
    k_end = tl.minimum(n_lo + BLOCK_N, N)

    k = k_start
    while k < k_end:
        kk = k + offs_k
        a = tl.load(A_ptr + offs_m[:, None] * N + kk[None, :],
                    mask=(offs_m[:, None] < N) & (kk[None, :] < N), other=0.0)
        b = tl.load(B_ptr + kk[:, None] * N + offs_n[None, :],
                    mask=(kk[:, None] < N) & (offs_n[None, :] < N), other=0.0)
        acc += tl.dot(a, b, allow_tf32=False)
        k += BLOCK_K

    mask_tri = offs_m[:, None] <= offs_n[None, :]
    acc = tl.where(mask_tri, acc, 0.0)
    store_mask = (offs_m[:, None] < N) & (offs_n[None, :] < N)
    tl.store(C_ptr + offs_m[:, None] * N + offs_n[None, :], acc, mask=store_mask)


class ModelNew(nn.Module):
    def __init__(self):
        super().__init__()

    def forward(self, A, B):
        A = A.contiguous()
        B = B.contiguous()
        N = A.shape[0]
        C = torch.zeros_like(A)
        BLOCK_M = 64
        BLOCK_N = 64
        BLOCK_K = 32
        grid = (triton.cdiv(N, BLOCK_M), triton.cdiv(N, BLOCK_N))
        triu_mm_kernel[grid](A, B, C, N, BLOCK_M=BLOCK_M, BLOCK_N=BLOCK_N, BLOCK_K=BLOCK_K,
                             num_warps=4, num_stages=3)
        return C
```