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

### A verified kernel for a related network problem (L3_021 EfficientNetMBConv, 1.05x): 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 bn_relu6_kernel(x_ptr, out_ptr, scale_ptr, shift_ptr, HW, C, n_elem, BLOCK: tl.constexpr):
    pid = tl.program_id(0)
    offs = pid * BLOCK + tl.arange(0, BLOCK)
    mask = offs < n_elem
    c = (offs // HW) % C
    x = tl.load(x_ptr + offs, mask=mask, other=0.0)
    s = tl.load(scale_ptr + c, mask=mask, other=0.0)
    sh = tl.load(shift_ptr + c, mask=mask, other=0.0)
    y = x * s + sh
    y = tl.minimum(tl.maximum(y, 0.0), 6.0)
    tl.store(out_ptr + offs, y, mask=mask)

@triton.jit
def bn_kernel(x_ptr, out_ptr, scale_ptr, shift_ptr, HW, C, n_elem, BLOCK: tl.constexpr):
    pid = tl.program_id(0)
    offs = pid * BLOCK + tl.arange(0, BLOCK)
    mask = offs < n_elem
    c = (offs // HW) % C
    x = tl.load(x_ptr + offs, mask=mask, other=0.0)
    s = tl.load(scale_ptr + c, mask=mask, other=0.0)
    sh = tl.load(shift_ptr + c, mask=mask, other=0.0)
    y = x * s + sh
    tl.store(out_ptr + offs, y, mask=mask)

@triton.jit
def bn_add_kernel(x_ptr, id_ptr, out_ptr, scale_ptr, shift_ptr, HW, C, n_elem, BLOCK: tl.constexpr):
    pid = tl.program_id(0)
    offs = pid * BLOCK + tl.arange(0, BLOCK)
    mask = offs < n_elem
    c = (offs // HW) % C
    x = tl.load(x_ptr + offs, mask=mask, other=0.0)
    idv = tl.load(id_ptr + offs, mask=mask, other=0.0)
    s = tl.load(scale_ptr + c, mask=mask, other=0.0)
    sh = tl.load(shift_ptr + c, mask=mask, other=0.0)
    y = x * s + sh + idv
    tl.store(out_ptr + offs, y, mask=mask)

def apply_bn_relu6(x, scale, shift):
    N, C, H, W = x.shape
    out = torch.empty_like(x)
    HW = H * W
    n = x.numel()
    grid = (triton.cdiv(n, 1024),)
    bn_relu6_kernel[grid](x, out, scale.contiguous(), shift.contiguous(), HW, C, n, BLOCK=1024)
    return out

def apply_bn(x, scale, shift):
    N, C, H, W = x.shape
    out = torch.empty_like(x)
    HW = H * W
    n = x.numel()
    grid = (triton.cdiv(n, 1024),)
    bn_kernel[grid](x, out, scale.contiguous(), shift.contiguous(), HW, C, n, BLOCK=1024)
    return out

def apply_bn_add(x, identity, scale, shift):
    N, C, H, W = x.shape
    out = torch.empty_like(x)
    HW = H * W
    n = x.numel()
    grid = (triton.cdiv(n, 1024),)
    bn_add_kernel[grid](x, identity.contiguous(), out, scale.contiguous(), shift.contiguous(), HW, C, n, BLOCK=1024)
    return out

class ModelNew(nn.Module):
    def __init__(self, in_channels, out_channels, kernel_size, stride, expand_ratio):
        super(ModelNew, self).__init__()
        self.use_residual = (stride == 1 and in_channels == out_channels)
        hidden_dim = in_channels * expand_ratio
        self.expand_ratio = expand_ratio

        if expand_ratio != 1:
            self.expand_conv = nn.Sequential(
                nn.Conv2d(in_channels, hidden_dim, kernel_size=1, stride=1, padding=0, bias=False),
                nn.BatchNorm2d(hidden_dim),
                nn.ReLU6(inplace=True)
            )

        self.depthwise_conv = nn.Sequential(
            nn.Conv2d(hidden_dim, hidden_dim, kernel_size=kernel_size, stride=stride, padding=(kernel_size-1)//2, groups=hidden_dim, bias=False),
            nn.BatchNorm2d(hidden_dim),
            nn.ReLU6(inplace=True)
        )

        self.project_conv = nn.Sequential(
            nn.Conv2d(hidden_dim, out_channels, kernel_size=1, stride=1, padding=0, bias=False),
            nn.BatchNorm2d(out_channels)
        )

    def _bn_params(self, bn, x):
        if self.training:
            mean = x.mean(dim=[0, 2, 3])
            var = x.var(dim=[0, 2, 3], unbiased=False)
            n = x.numel() // x.shape[1]
            with torch.no_grad():
                bn.running_mean.mul_(1 - bn.momentum).add_(bn.momentum * mean)
                bn.running_var.mul_(1 - bn.momentum).add_(bn.momentum * var * n / (n - 1))
        else:
            mean = bn.running_mean
            var = bn.running_var
        scale = bn.weight / torch.sqrt(var + bn.eps)
        shift = bn.bias - mean * scale
        return scale, shift

    def forward(self, x):
        identity = x

        if hasattr(self, 'expand_conv'):
            x = self.expand_conv[0](x)
            s, sh = self._bn_params(self.expand_conv[1], x)
            x = apply_bn_relu6(x, s, sh)

        x = self.depthwise_conv[0](x)
        s, sh = self._bn_params(self.depthwise_conv[1], x)
        x = apply_bn_relu6(x, s, sh)

        x = self.project_conv[0](x)
        s, sh = self._bn_params(self.project_conv[1], x)
        if self.use_residual:
            x = apply_bn_add(x, identity, s, sh)
        else:
            x = apply_bn(x, s, sh)

        return x
```

### A verified kernel for a related network problem (L3_027 RegNet, 1.00x): adapt what applies
```python
import torch
import torch.nn as nn
import triton
import triton.language as tl

@triton.jit
def bn_relu_kernel(x_ptr, out_ptr, scale_ptr, shift_ptr, HW, C, n_elem, BLOCK: tl.constexpr):
    pid = tl.program_id(0)
    offs = pid * BLOCK + tl.arange(0, BLOCK)
    mask = offs < n_elem
    c = (offs // HW) % C
    x = tl.load(x_ptr + offs, mask=mask, other=0.0)
    s = tl.load(scale_ptr + c, mask=mask, other=0.0)
    sh = tl.load(shift_ptr + c, mask=mask, other=0.0)
    y = x * s + sh
    y = tl.maximum(y, 0.0)
    tl.store(out_ptr + offs, y, mask=mask)

def apply_bn_relu(x, scale, shift):
    N, C, H, W = x.shape
    out = torch.empty_like(x)
    HW = H * W
    n = x.numel()
    grid = (triton.cdiv(n, 1024),)
    bn_relu_kernel[grid](x, out, scale.contiguous(), shift.contiguous(), HW, C, n, BLOCK=1024)
    return out

class ModelNew(nn.Module):
    def __init__(self, input_channels, stages, block_widths, output_classes):
        super(ModelNew, self).__init__()
        self.stages = stages
        self.block_widths = block_widths
        layers = []
        current = input_channels
        for i in range(stages):
            layers.append(nn.Sequential(
                nn.Conv2d(current, block_widths[i], kernel_size=3, padding=1),
                nn.BatchNorm2d(block_widths[i]),
                nn.ReLU(),
                nn.Conv2d(block_widths[i], block_widths[i], kernel_size=3, padding=1),
                nn.BatchNorm2d(block_widths[i]),
                nn.ReLU(),
                nn.MaxPool2d(kernel_size=2, stride=2),
            ))
            current = block_widths[i]
        self.feature_extractor = nn.Sequential(*layers)
        self.fc = nn.Linear(block_widths[-1], output_classes)

    def _bn_params(self, bn, x):
        if self.training:
            mean = x.mean(dim=[0, 2, 3])
            var = x.var(dim=[0, 2, 3], unbiased=False)
            n = x.numel() // x.shape[1]
            with torch.no_grad():
                bn.running_mean.mul_(1 - bn.momentum).add_(bn.momentum * mean)
                bn.running_var.mul_(1 - bn.momentum).add_(bn.momentum * var * n / (n - 1))
        else:
            mean = bn.running_mean
            var = bn.running_var
        scale = bn.weight / torch.sqrt(var + bn.eps)
        shift = bn.bias - mean * scale
        return scale, shift

    def forward(self, x):
        for stage in self.feature_extractor:
            conv1 = stage[0]
            bn1 = stage[1]
            conv2 = stage[3]
            bn2 = stage[4]
            pool = stage[6]
            x = conv1(x)
            s, sh = self._bn_params(bn1, x)
            x = apply_bn_relu(x, s, sh)
            x = conv2(x)
            s, sh = self._bn_params(bn2, x)
            x = apply_bn_relu(x, s, sh)
            x = pool(x)
        x = torch.mean(x, dim=[2, 3])
        x = self.fc(x)
        return x
```