## 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_063 Gemm_ReLU_Divide, 1.01x): adapt what applies
```python
import torch
import torch.nn as nn
import triton
import triton.language as tl

@triton.jit
def epilogue_kernel(x_ptr, bias_ptr, out_ptr, inv_div, M, N, BLOCK: tl.constexpr):
    pid = tl.program_id(0)
    offs = pid * BLOCK + tl.arange(0, BLOCK)
    mask = offs < M * N
    col = offs % N
    v = tl.load(x_ptr + offs, mask=mask, other=0.0)
    b = tl.load(bias_ptr + col, mask=mask, other=0.0)
    y = tl.maximum(v + b, 0.0) * inv_div
    tl.store(out_ptr + offs, y, mask=mask)

class ModelNew(nn.Module):
    def __init__(self, in_features, out_features, divisor):
        super().__init__()
        self.linear = nn.Linear(in_features, out_features)
        self.divisor = divisor

    def forward(self, x):
        z = torch.mm(x, self.linear.weight.t())
        M, N = z.shape
        out = torch.empty_like(z)
        BLOCK = 1024
        grid = (triton.cdiv(M * N, BLOCK),)
        epilogue_kernel[grid](z, self.linear.bias, out, 1.0 / self.divisor, M, N, BLOCK=BLOCK)
        return out
```

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

@triton.jit
def stats_kernel(x_ptr, mean_ptr, var_ptr, M, N, BLOCK_M: tl.constexpr):
    col = tl.program_id(0)
    offs = tl.arange(0, BLOCK_M)
    acc = tl.zeros((), dtype=tl.float32)
    acc2 = tl.zeros((), dtype=tl.float32)
    for start in range(0, M, BLOCK_M):
        idx = start + offs
        mask = idx < M
        v = tl.load(x_ptr + idx * N + col, mask=mask, other=0.0)
        acc += tl.sum(v)
        acc2 += tl.sum(v * v)
    mean = acc / M
    var = acc2 / M - mean * mean
    tl.store(mean_ptr + col, mean)
    tl.store(var_ptr + col, var)

@triton.jit
def apply_kernel(x_ptr, out_ptr, mean_ptr, var_ptr, gamma_ptr, beta_ptr, bias_ptr,
                 eps, divide_value, M, N, BLOCK: tl.constexpr):
    pid = tl.program_id(0)
    offs = pid * BLOCK + tl.arange(0, BLOCK)
    mask = offs < M * N
    col = offs % N
    v = tl.load(x_ptr + offs, mask=mask, other=0.0)
    mean = tl.load(mean_ptr + col, mask=mask, other=0.0)
    var = tl.load(var_ptr + col, mask=mask, other=0.0)
    gamma = tl.load(gamma_ptr + col, mask=mask, other=0.0)
    beta = tl.load(beta_ptr + col, mask=mask, other=0.0)
    bias = tl.load(bias_ptr)
    inv = 1.0 / tl.sqrt(var + eps)
    y = (v - mean) * inv * gamma + beta
    y = (y + bias) / divide_value
    y = y * (1.0 / (1.0 + tl.exp(-y)))
    tl.store(out_ptr + offs, y, mask=mask)

class ModelNew(nn.Module):
    def __init__(self, in_features, out_features, bn_eps=1e-5, bn_momentum=0.1, bias_shape=(1,), divide_value=1.0):
        super().__init__()
        self.matmul = nn.Linear(in_features, out_features)
        self.bn = nn.BatchNorm1d(out_features, eps=bn_eps, momentum=bn_momentum)
        self.bias = nn.Parameter(torch.randn(bias_shape))
        self.divide_value = divide_value
        self.bn_eps = bn_eps
        self.bn_momentum = bn_momentum

    def forward(self, x):
        z = self.matmul(x).contiguous()
        M, N = z.shape
        mean = torch.empty(N, device=z.device, dtype=torch.float32)
        var = torch.empty(N, device=z.device, dtype=torch.float32)
        stats_kernel[(N,)](z, mean, var, M, N, BLOCK_M=1024)

        with torch.no_grad():
            self.bn.running_mean.mul_(1 - self.bn_momentum).add_(self.bn_momentum * mean)
            unbiased = var * (M / (M - 1))
            self.bn.running_var.mul_(1 - self.bn_momentum).add_(self.bn_momentum * unbiased)
            self.bn.num_batches_tracked.add_(1)

        out = torch.empty_like(z)
        bias_flat = self.bias.reshape(-1)
        BLOCK = 1024
        grid = (triton.cdiv(M * N, BLOCK),)
        apply_kernel[grid](z, out, mean, var, self.bn.weight, self.bn.bias, bias_flat,
                           self.bn_eps, self.divide_value, M, N, BLOCK=BLOCK)
        return out
```