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


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


### Verified technique (E-H6, admitted by the KFORGE engine): Attention and batched matmuls on bf16x3: a flash-attention Triton kernel with in-register hi/lo split for QK^T and PV
For attention blocks and attention layers inside networks (ViT/Swin/transformer, head_dim 32-128), a FlashAttention-2-style Triton kernel should beat eager fp32 matmul -> softmax -> matmul by 2-4x. The kernel computes S = QK^T and O = PV with three bf16 tl.dot calls each, keeps online softmax in fp32, and never materializes the [L, L] scores.
Physics: The eager reference does two SIMT fp32 bmm calls plus a softmax that writes and reads an L x L fp32 tensor, which is HBM-bound at large L. The fused kernel removes that traffic and moves the flops to bf16 mma. Reduction lengths are tiny (K = head_dim for QK^T, block length for PV with fp32 rescaling), so the 3-product error sits far below 1e-4. Softmax exponent error stays around 1e-6 because the score error is about 2^-16 relative.
How: Use grid = (cdiv(L, BLOCK_M), B*H) with BLOCK_M = 64-128, BLOCK_N = 64, and head_dim held as a constexpr. Split the Q tile into hi/lo bf16 once per program and split each K/V tile on load. Compute s = (dot(qh, kl) + dot(ql, kh) + dot(qh, kh)) * scale, apply masks if present, update the running max m and sum l in fp32, and rescale acc. Split p = exp(s - m) into hi/lo and compute acc += dot(pl, vh) + dot(ph, vl) + dot(ph, vh). Write O/l in the layout the reference returns, handling the transpose and reshape via strides with no .contiguous(). If the reference adds a relative-position bias (as in Swin), cache the parameter-only bias tensor in self.<x> keyed on data_ptr and load it inside the kernel. Fall back to P1 torch.mm for non-attention bmm with a shared weight by reshaping to 2D.
- Verified on training problems (protected evaluator): L3_043 MinGPTCausalAttention: 4.16x vs 3.74x with the previous library.
- Verified example implementing it (L3_043, 4.16x; adapt what applies):
```python
import math
import torch
import torch.nn as nn
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):
    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


@triton.autotune(
    configs=[
        triton.Config({"BM": 64, "BN": 64}, num_stages=2, num_warps=4),
        triton.Config({"BM": 128, "BN": 64}, num_stages=2, num_warps=8),
        triton.Config({"BM": 64, "BN": 32}, num_stages=2, num_warps=4),
        triton.Config({"BM": 128, "BN": 32}, num_stages=2, num_warps=8),
    ],
    key=["T", "C"],
)
@triton.jit
def attn_kernel(qkv, bias, out, T, H, C, scale_log2,
                HS: tl.constexpr, D: tl.constexpr, BM: tl.constexpr, BN: tl.constexpr):
    num_m = tl.cdiv(T, BM)
    pid_m = num_m - 1 - tl.program_id(0)
    bh = tl.program_id(1)
    b = bh // H
    h = bh % H
    offs_m = pid_m * BM + tl.arange(0, BM)
    offs_n0 = tl.arange(0, BN)
    offs_d = tl.arange(0, D)
    dmask = offs_d < HS
    rs = 3 * C
    base = qkv + b.to(tl.int64) * T * rs + h * HS
    mrow = offs_m < T
    q = tl.load(base + offs_m[:, None] * rs + offs_d[None, :],
                mask=mrow[:, None] & dmask[None, :], other=0.0)
    bq = tl.load(bias + h * HS + offs_d, mask=dmask, other=0.0)
    bk = tl.load(bias + C + h * HS + offs_d, mask=dmask, other=0.0)
    bv = tl.load(bias + 2 * C + h * HS + offs_d, mask=dmask, other=0.0)
    q = q + bq[None, :]
    qh = q.to(tl.bfloat16)
    ql = (q - qh.to(tl.float32)).to(tl.bfloat16)
    m_i = tl.full((BM,), float("-inf"), tl.float32)
    l_i = tl.zeros((BM,), tl.float32)
    acc = tl.zeros((BM, D), tl.float32)
    hi = tl.minimum((pid_m + 1) * BM, T)
    for start_n in range(0, hi, BN):
        offs_n = start_n + offs_n0
        nmask = offs_n < T
        kT = tl.load(base + C + offs_n[None, :] * rs + offs_d[:, None],
                     mask=nmask[None, :] & dmask[:, None], other=0.0) + bk[:, None]
        kh = kT.to(tl.bfloat16)
        kl = (kT - kh.to(tl.float32)).to(tl.bfloat16)
        s2 = tl.dot(qh, kl)
        s2 = tl.dot(ql, kh, s2)
        s = tl.dot(qh, kh)
        s = (s + s2) * scale_log2
        valid = (offs_m[:, None] >= offs_n[None, :]) & nmask[None, :]
        s = tl.where(valid, s, float("-inf"))
        m_new = tl.maximum(m_i, tl.max(s, 1))
        alpha = tl.exp2(m_i - m_new)
        p = tl.exp2(s - m_new[:, None])
        l_i = l_i * alpha + tl.sum(p, 1)
        acc = acc * alpha[:, None]
        v = tl.load(base + 2 * C + offs_n[:, None] * rs + offs_d[None, :],
                    mask=nmask[:, None] & dmask[None, :], other=0.0) + bv[None, :]
        vh = v.to(tl.bfloat16)
        vl = (v - vh.to(tl.float32)).to(tl.bfloat16)
        ph = p.to(tl.bfloat16)
        pl = (p - ph.to(tl.float32)).to(tl.bfloat16)
        acc = tl.dot(pl, vh, acc)
        acc = tl.dot(ph, vl, acc)
        acc = tl.dot(ph, vh, acc)
        m_i = m_new
    o = acc / l_i[:, None]
    oh = o.to(tl.bfloat16)
    ol = (o - oh.to(tl.float32)).to(tl.bfloat16)
    obase = out + (b.to(tl.int64) * T + offs_m[:, None]) * rs + h * HS + offs_d[None, :]
    om = mrow[:, None] & dmask[None, :]
    tl.store(obase, oh, mask=om)
    tl.store(obase + C, oh, mask=om)
    tl.store(obase + 2 * C, ol, mask=om)


class ModelNew(nn.Module):
    def __init__(self, n_embd, n_head, attn_pdrop, resid_pdrop, max_seqlen):
        super().__init__()
        assert n_embd % n_head == 0
        self.c_attn = nn.Linear(n_embd, 3 * n_embd)
        self.c_proj = nn.Linear(n_embd, n_embd)
        self.attn_dropout = nn.Dropout(attn_pdrop)
        self.resid_dropout = nn.Dropout(resid_pdrop)
        self.register_buffer("bias", torch.tril(torch.ones(max_seqlen, max_seqlen))
                             .view(1, 1, max_seqlen, max_seqlen))
        self.n_head = n_head
        self.n_embd = n_embd

    def forward(self, x):
        B, T, C = x.size()
        H = self.n_head
        hs = C // H
        x2 = x.reshape(B * T, C)
        qkv = torch.mm(split3(x2, 0), split3(self.c_attn.weight, 1).t(), out_dtype=torch.float32)
        ysplit = torch.empty((B * T, 3 * C), device=x.device, dtype=torch.bfloat16)
        D = triton.next_power_of_2(hs)
        scale_log2 = (1.0 / math.sqrt(hs)) * 1.4426950408889634
        grid = lambda META: (triton.cdiv(T, META["BM"]), B * H)
        attn_kernel[grid](qkv, self.c_attn.bias, ysplit, T, H, C, scale_log2, HS=hs, D=D)
        y = torch.mm(ysplit, split3(self.c_proj.weight, 1).t(), out_dtype=torch.float32)
        y.add_(self.c_proj.bias)
        return y.view(B, T, C)
```


## 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 attention problems (wrong results or slower than eager)
- Pitfall: Avoid per-row Triton LayerNorm with one program per row and huge num_warps for short rows; launch overhead plus extra .contiguous() copies make it slower than PyTorch's fused layer_norm.
- Pitfall: In large transformer models, replacing only LayerNorm (or fused LN+residual) with custom kernels while matmuls, attention, and window permutes stay eager yields slowdown; fuse the attention core or the dominant ops instead.
- Pitfall: Avoid extra .contiguous()/reshape copies and per-call Python wrapper overhead around small Triton launches; on small tensors launch latency outweighs the gain over cuDNN/ATen fused kernels.
- Pitfall: Don'...'t beat eager.
- Pitfall: Handle every optional forward argument (e.g. initial_states) and every intermediate the reference includes, even when it defaults to zeros. Verify the returned tensor is the same slice and quantity the reference returns.
- Pitfall: Do not algebraically restructure the reference'...'s exact operation order and test against it before optimizing.