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

### Verified technique (E-H4obs, observed in a verified solution; admitted by the KFORGE engine): bf16x3 split-precision GEMM via K-concatenation on tensor cores
Each fp32 input is split into a bf16 "hi" part plus a bf16 "lo" residual (lo = x − hi). The pieces are laid out along the K dimension so that one ordinary bf16 GEMM computes the cross terms hi·hi + hi·lo + lo·hi in a single pass. A's rows become [hi | hi | lo] and B's columns become [hi ; lo ; hi]. K is also cut into 4 chunks run as one batched bmm with fp32 output. A small elementwise kernel then sums the chunk partials, which shortens each accumulation chain and adds parallelism (split-K).

Numerically, hi carries about 8 mantissa bits and lo the next 8, so the result has roughly 16 bits of effective precision. The dropped lo·lo term is about 2^-16 relative, and accumulation stays in fp32. That lands near 1e-5 relative error, inside a 1e-4 tolerance, while running on bf16 tensor cores (about 312 TFLOPS on A100). The alternative is fp32 FMA (about 19.5 TFLOPS) when TF32 is disallowed. Since bf16 shares fp32's exponent range, the split does not overflow.

The technique applies to any large, compute-bound fp32 matmul, conv-as-GEMM, or batched GEMM where strict fp32 accuracy is checked but TF32 is unavailable or too imprecise. Things to watch:
- **Cost:** it does 3× the FLOPs and writes 3× bf16 copies of the operands, so it loses on small or memory-bound shapes.
- **Tiny residuals:** lo values can underflow or flush to zero for very small or denormal inputs.
- **Error growth:** error grows with K and with ill-conditioned or cancellation-heavy data, so validate on worst-case magnitudes.
- **Dependency:** `torch.bmm(..., out_dtype=float32)` requires a recent PyTorch.
- Verified on a training problem (protected evaluator): L1_002 Standard_matrix_multiplication_: 3.56x vs 3.25x with the previous library.
- Verified example (L1_002, 3.56x; adapt what applies):
```python
import torch
import torch.nn as nn
import triton
import triton.language as tl


@triton.jit
def split_rows_kernel(src, dst, R, K, KC, BR: tl.constexpr, BKK: tl.constexpr):
    # src [R, K] fp32 -> dst [C, R, 3KC] bf16, each chunk: [hi | hi | lo]
    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)
    c = k // KC
    kk = k - c * KC
    base = dst + c[None, :] * (R * 3 * KC) + r[:, None] * (3 * KC) + kk[None, :]
    tl.store(base, h, mask=m)
    tl.store(base + KC, h, mask=m)
    tl.store(base + 2 * KC, lo, mask=m)


@triton.jit
def split_blocks_kernel(src, dst, T, N, KC, BLOCK: tl.constexpr):
    # src [K, N] fp32 flat -> dst [C, 3KC, N] bf16, each chunk: [hi; lo; hi]
    i = tl.program_id(0) * BLOCK + tl.arange(0, BLOCK)
    m = i < T
    v = tl.load(src + i, mask=m, other=0.0)
    h = v.to(tl.bfloat16)
    lo = (v - h.to(tl.float32)).to(tl.bfloat16)
    k = i // N
    n = i - k * N
    c = k // KC
    kk = k - c * KC
    o = c * (3 * KC * N) + kk * N + n
    tl.store(dst + o, h, mask=m)
    tl.store(dst + o + KC * N, lo, mask=m)
    tl.store(dst + o + 2 * KC * N, h, mask=m)


@triton.jit
def chunk_sum_kernel(src, dst, T, C: tl.constexpr, BLOCK: tl.constexpr):
    i = tl.program_id(0) * BLOCK + tl.arange(0, BLOCK)
    m = i < T
    acc = tl.load(src + i, mask=m, other=0.0)
    for c in tl.static_range(1, C):
        acc += tl.load(src + c * T + i, mask=m, other=0.0)
    tl.store(dst + i, acc, mask=m)


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

    def forward(self, A: torch.Tensor, B: torch.Tensor) -> torch.Tensor:
        A = A.contiguous()
        B = B.contiguous()
        M, K = A.shape
        N = B.shape[1]
        C = 4 if K % 4 == 0 else 1
        KC = K // C
        a3 = torch.empty((C, M, 3 * KC), device=A.device, dtype=torch.bfloat16)
        split_rows_kernel[(triton.cdiv(M, 32), triton.cdiv(K, 128))](A, a3, M, K, KC, BR=32, BKK=128)
        b3 = torch.empty((C, 3 * KC, N), device=B.device, dtype=torch.bfloat16)
        T = K * N
        split_blocks_kernel[(triton.cdiv(T, 2048),)](B, b3, T, N, KC, BLOCK=2048)
        part = torch.bmm(a3, b3, out_dtype=torch.float32)
        if C == 1:
            return part[0]
        out = torch.empty((M, N), device=A.device, dtype=torch.float32)
        TO = M * N
        chunk_sum_kernel[(triton.cdiv(TO, 2048),)](part, out, TO, C=C, BLOCK=2048)
        return out
```


## 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
- Pitfall: Output mismatch beyond 1e-4 occurred repeatedly. Fix by forcing exact fp32 tl.dot(allow_tf32=False) and exact masking: K bounds at block granularity, tril mask, and zero-initialized output for skipped tiles.
- 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.

### 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: Strict fp32 Triton tl.dot (allow_tf32=False) rarely beats cuBLAS on matmul; a hand split-K version was correct but still 0.88x. Fuse surrounding ops or exploit structure instead of reimplementing plain GEMM.
- 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.

### A verified kernel for a related matmul problem (L2_051 Gemm_Subtract_GlobalAvgPool_LogSumExp_GELU_ResidualAdd, 36.02x): adapt what applies
```python
import torch
import torch.nn as nn
import torch.nn.functional as F
import triton
import triton.language as tl


@triton.jit
def add_kernel(x_ptr, s_ptr, out_ptr, n_feat, n, BLOCK: tl.constexpr):
    pid = tl.program_id(0)
    offs = pid * BLOCK + tl.arange(0, BLOCK)
    mask = offs < n
    row = offs // n_feat
    v = tl.load(x_ptr + offs, mask=mask)
    s = tl.load(s_ptr + row, mask=mask)
    tl.store(out_ptr + offs, v + s, mask=mask)


class ModelNew(nn.Module):
    def __init__(self, in_features, out_features, bias=True):
        super().__init__()
        self.gemm = nn.Linear(in_features, out_features, bias=bias)
        self.subtract = nn.Parameter(torch.randn(out_features))
        self.out_features = out_features

    def forward(self, x):
        # mean over output features of (Wx + b - subtract) = (colsum/out)·x + c
        colsum = self.gemm.weight.sum(dim=0) / self.out_features  # (in,)
        if self.gemm.bias is not None:
            c = (self.gemm.bias - self.subtract).mean()
        else:
            c = (-self.subtract).mean()
        s = x @ colsum + c  # (batch,)
        # logsumexp over single element (dim=1 after keepdim mean) is identity
        s = F.gelu(s)
        out = torch.empty_like(x)
        n = x.numel()
        n_feat = x.shape[1]
        BLOCK = 1024
        grid = (triton.cdiv(n, BLOCK),)
        add_kernel[grid](x, s, out, n_feat, n, BLOCK=BLOCK)
        return out
```

### 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
```