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