```python
# Copyright (c) Meta Platforms, Inc. and affiliates.
# All rights reserved.
#
# This source code is licensed under the license found in the
# LICENSE file in the root directory of this source tree.

import torch
import torch.nn as nn
import torch.nn.functional as F
from torch.utils.cpp_extension import load_inline

torch.backends.cudnn.benchmark = True
torch.backends.cudnn.allow_tf32 = True

# Custom CUDA pipeline (channels_last):
#   0) nchw_to_nhwc_kernel       : smem-tiled transpose of the (small) input so
#                                  cuDNN runs its native NHWC tensor-core conv.
#                                  Now uses fully-coalesced float2 writes.
#   1) stats_nhwc_kernel         : fully-coalesced per-(batch, group) sum/sumsq
#                                  accumulation (float4 loads + __ldcs + warp
#                                  pair-combine + smem + atomics). Conv bias is
#                                  FUSED here (added per element) so the conv
#                                  runs with bias=None, eliminating cuDNN's
#                                  separate ~770us broadcast-bias kernel.
#   2) lse_nhwc_kernel           : fused (bias +) GroupNorm affine + Tanh +
#                                  HardSwish + residual add + LogSumExp over
#                                  channels, two-phase warp reduction, float4
#                                  streaming stores, split full/tail loops.
#   (+ generic fallback kernels for uncommon channel configs)
_cuda_src = r"""
#include <torch/extension.h>
#include <ATen/cuda/CUDAContext.h>
#include <cuda_runtime.h>

#define THREADS 256
#define NEG_INF __int_as_float(0xff800000)
#define MAX_CPL 4

__device__ __forceinline__ float fast_tanh(float y) {
    // tanh(y) = sign(y) * (1 - e^{-2|y|}) / (1 + e^{-2|y|}); no overflow.
    float e = __expf(-2.0f * fabsf(y));
    float t = __fdividef(1.0f - e, 1.0f + e);
    return copysignf(t, y);
}

__device__ __forceinline__ float act_res(float v, float a, float b) {
    // groupnorm affine: y = v * a + b, then tanh, then hardswish, then +v.
    float y = v * a + b;
    float t = fast_tanh(y);
    float hs = t * fminf(fmaxf(t + 3.0f, 0.0f), 6.0f) * (1.0f / 6.0f);
    return v + hs;
}

// ---------------------------------------------------------------------------
// Kernel 0: NCHW -> NHWC transpose, smem tiled. Each block handles 256 spatial
// positions x 8 channels of one image. Coalesced reads AND (float2) writes.
// ---------------------------------------------------------------------------
__global__ void nchw_to_nhwc_kernel(const float *__restrict__ in,
                                    float *__restrict__ out,
                                    int C, int HW) {
    __shared__ float sm[8 * 260];
    const int hw0 = blockIdx.x * 256;
    const int c0 = blockIdx.y * 8;
    const int n = blockIdx.z;
    const int nc = min(8, C - c0);
    const int t = threadIdx.x;
    const bool hv = (hw0 + t) < HW;
    const float *in_n = in + (size_t)n * (size_t)C * (size_t)HW;
    float vals[8];
    #pragma unroll
    for (int c = 0; c < 8; ++c) {
        vals[c] = (c < nc && hv) ? __ldg(in_n + (size_t)(c0 + c) * (size_t)HW + hw0 + t)
                                 : 0.0f;
    }
    #pragma unroll
    for (int c = 0; c < 8; ++c) {
        if (c < nc) sm[c * 260 + t] = vals[c];
    }
    __syncthreads();
    float *out_n = out + (size_t)n * (size_t)HW * (size_t)C + (size_t)hw0 * (size_t)C + c0;
    if ((C & 1) == 0 && (nc & 1) == 0) {
        const int cpairs = nc >> 1;
        const int total = 256 * cpairs;
        const int Crow2 = C >> 1;
        float2 *out2 = reinterpret_cast<float2 *>(out_n);
        for (int f = t; f < total; f += blockDim.x) {
            const int i = f / cpairs;
            const int cp = f - i * cpairs;
            if (hw0 + i < HW) {
                out2[(size_t)i * (size_t)Crow2 + cp] =
                    make_float2(sm[(2 * cp) * 260 + i], sm[(2 * cp + 1) * 260 + i]);
            }
        }
    } else {
        const int total = 256 * nc;
        for (int f = t; f < total; f += blockDim.x) {
            const int i = f / nc;
            const int c = f - i * nc;
            if (hw0 + i < HW) {
                out_n[(size_t)i * (size_t)C + c] = sm[c * 260 + i];
            }
        }
    }
}

// ---------------------------------------------------------------------------
// Kernel 1 (fast): per-(n, group) sum / sumsq of the NHWC conv output, with
// the conv bias fused in. Requires: C % 4 == 0, CperG % 4 == 0,
// (C/4) <= THREADS, THREADS % (C/4) == 0.
// ---------------------------------------------------------------------------
__global__ void stats_nhwc_kernel(const float4 *__restrict__ x4,
                                  const float4 *__restrict__ cb4,
                                  float *__restrict__ acc,
                                  int HW, int C, int G, int CperG, int chunk) {
    const int n = blockIdx.y;
    const int hw0 = blockIdx.x * chunk;
    const int span = min(chunk, HW - hw0);
    if (span <= 0) return;
    const int f4perHw = C >> 2;
    const int f4idx = threadIdx.x % f4perHw;
    const int hwSub = threadIdx.x / f4perHw;
    const int hwStep = blockDim.x / f4perHw;
    const int g = (f4idx << 2) / CperG;
    const float4 bv = __ldg(cb4 + f4idx);
    const float4 *base = x4 + ((size_t)n * (size_t)HW + hw0) * (size_t)f4perHw;
    float sum = 0.0f, sq = 0.0f;
    #pragma unroll 4
    for (int hw = hwSub; hw < span; hw += hwStep) {
        float4 v = __ldcs(base + (size_t)hw * (size_t)f4perHw + f4idx);
        v.x += bv.x; v.y += bv.y; v.z += bv.z; v.w += bv.w;
        sum += (v.x + v.y) + (v.z + v.w);
        sq  += (v.x * v.x + v.y * v.y) + (v.z * v.z + v.w * v.w);
    }
    // Lanes i and i+16 see the same channels when f4perHw divides 16: combine
    // pairs first to halve the shared-memory atomic traffic.
    if (f4perHw <= 16) {
        sum += __shfl_xor_sync(0xffffffffu, sum, 16);
        sq  += __shfl_xor_sync(0xffffffffu, sq, 16);
    }
    __shared__ float sm[64 * 2];
    for (int i = threadIdx.x; i < 2 * G; i += blockDim.x) sm[i] = 0.0f;
    __syncthreads();
    if (f4perHw > 16 || (threadIdx.x & 16) == 0) {
        atomicAdd(&sm[g * 2 + 0], sum);
        atomicAdd(&sm[g * 2 + 1], sq);
    }
    __syncthreads();
    float *accg = acc + (size_t)n * (size_t)G * 2;
    for (int i = threadIdx.x; i < G; i += blockDim.x) {
        atomicAdd(&accg[i * 2 + 0], sm[i * 2 + 0]);
        atomicAdd(&accg[i * 2 + 1], sm[i * 2 + 1]);
    }
}

// ---------------------------------------------------------------------------
// Kernel 1 (generic fallback): one block per (n, group), scalar strided loop.
// ---------------------------------------------------------------------------
__global__ void stats_nhwc_generic_kernel(const float *__restrict__ x,
                                          const float *__restrict__ cb,
                                          float *__restrict__ acc,
                                          int HW, int C, int G, int CperG) {
    const int ng = blockIdx.x;
    const int n = ng / G;
    const int g = ng - n * G;
    const int c0 = g * CperG;
    const float *xb = x + (size_t)n * (size_t)HW * (size_t)C;
    const int total = CperG * HW;
    float sum = 0.0f, sq = 0.0f;
    for (int i = threadIdx.x; i < total; i += blockDim.x) {
        const int hw = i / CperG;
        const int c = c0 + (i - hw * CperG);
        const float v = xb[(size_t)hw * (size_t)C + c] + __ldg(cb + c);
        sum += v;
        sq += v * v;
    }
    __shared__ float ssum[THREADS / 32];
    __shared__ float sssq[THREADS / 32];
    const int lane = threadIdx.x & 31;
    const int wid = threadIdx.x >> 5;
    #pragma unroll
    for (int o = 16; o > 0; o >>= 1) {
        sum += __shfl_down_sync(0xffffffffu, sum, o);
        sq += __shfl_down_sync(0xffffffffu, sq, o);
    }
    if (lane == 0) { ssum[wid] = sum; sssq[wid] = sq; }
    __syncthreads();
    if (wid == 0) {
        const int nw = THREADS / 32;
        sum = (lane < nw) ? ssum[lane] : 0.0f;
        sq = (lane < nw) ? sssq[lane] : 0.0f;
        #pragma unroll
        for (int o = 16; o > 0; o >>= 1) {
            sum += __shfl_down_sync(0xffffffffu, sum, o);
            sq += __shfl_down_sync(0xffffffffu, sq, o);
        }
        if (lane == 0) {
            acc[ng * 2 + 0] = sum;
            acc[ng * 2 + 1] = sq;
        }
    }
}

// ---------------------------------------------------------------------------
// Fused activation + two-phase LogSumExp for 4 spatial positions held in
// registers. ALL lanes execute the butterflies (shuffle participation).
// ---------------------------------------------------------------------------
__device__ __forceinline__ void quad_lse(float2 v0, float2 v1, float2 v2, float2 v3,
                                         float a0, float b0, float a1, float b1,
                                         float bs0, float bs1, bool active,
                                         float &m0, float &m1, float &m2, float &m3,
                                         float &s0, float &s1, float &s2, float &s3) {
    float r0 = 0.f, r1 = 0.f, r2 = 0.f, r3 = 0.f;
    float r4 = 0.f, r5 = 0.f, r6 = 0.f, r7 = 0.f;
    if (active) {
        r0 = act_res(v0.x + bs0, a0, b0); r1 = act_res(v0.y + bs1, a1, b1);
        r2 = act_res(v1.x + bs0, a0, b0); r3 = act_res(v1.y + bs1, a1, b1);
        r4 = act_res(v2.x + bs0, a0, b0); r5 = act_res(v2.y + bs1, a1, b1);
        r6 = act_res(v3.x + bs0, a0, b0); r7 = act_res(v3.y + bs1, a1, b1);
        m0 = fmaxf(r0, r1); m1 = fmaxf(r2, r3);
        m2 = fmaxf(r4, r5); m3 = fmaxf(r6, r7);
    } else {
        m0 = m1 = m2 = m3 = NEG_INF;
    }
    #pragma unroll
    for (int o = 16; o > 0; o >>= 1) {
        m0 = fmaxf(m0, __shfl_xor_sync(0xffffffffu, m0, o));
        m1 = fmaxf(m1, __shfl_xor_sync(0xffffffffu, m1, o));
        m2 = fmaxf(m2, __shfl_xor_sync(0xffffffffu, m2, o));
        m3 = fmaxf(m3, __shfl_xor_sync(0xffffffffu, m3, o));
    }
    s0 = s1 = s2 = s3 = 0.f;
    if (active) {
        s0 = __expf(r0 - m0) + __expf(r1 - m0);
        s1 = __expf(r2 - m1) + __expf(r3 - m1);
        s2 = __expf(r4 - m2) + __expf(r5 - m2);
        s3 = __expf(r6 - m3) + __expf(r7 - m3);
    }
    #pragma unroll
    for (int o = 16; o > 0; o >>= 1) {
        s0 += __shfl_xor_sync(0xffffffffu, s0, o);
        s1 += __shfl_xor_sync(0xffffffffu, s1, o);
        s2 += __shfl_xor_sync(0xffffffffu, s2, o);
        s3 += __shfl_xor_sync(0xffffffffu, s3, o);
    }
}

// ---------------------------------------------------------------------------
// Kernel 2 (fast): fused (bias +) GroupNorm + Tanh + HardSwish + residual +
// LogSumExp, NHWC. One warp handles 4 consecutive spatial positions; lane l
// owns channels (2l, 2l+1) via float2 loads. Requires: C even, C <= 64.
// ---------------------------------------------------------------------------
__global__ void lse_nhwc_kernel(const float *__restrict__ x,
                                const float *__restrict__ cb,
                                const float *__restrict__ w,
                                const float *__restrict__ b,
                                const float *__restrict__ acc,
                                float *__restrict__ out,
                                int HW, int C, int G, int CperG,
                                float inv_cnt, float eps) {
    const int n = blockIdx.y;
    const int lane = threadIdx.x & 31;
    const int wid = threadIdx.x >> 5;
    const int nwarps = blockDim.x >> 5;
    const int c0 = lane * 2;
    const int c1 = c0 + 1;
    const bool active = c0 < C;   // C even => c0 < C implies c1 < C
    float a0 = 0.f, b0 = 0.f, a1 = 0.f, b1 = 0.f, bs0 = 0.f, bs1 = 0.f;
    if (active) {
        const float *accn = acc + (size_t)n * (size_t)G * 2;
        const int g0 = c0 / CperG;
        const float sum0 = accn[g0 * 2 + 0], sq0 = accn[g0 * 2 + 1];
        const float mean0 = sum0 * inv_cnt;
        const float var0 = fmaxf(sq0 * inv_cnt - mean0 * mean0, 0.0f);
        const float rstd0 = rsqrtf(var0 + eps);
        a0 = w[c0] * rstd0;
        b0 = b[c0] - mean0 * a0;
        const int g1 = c1 / CperG;
        if (g1 == g0) {
            a1 = w[c1] * rstd0;
            b1 = b[c1] - mean0 * a1;
        } else {
            const float sum1 = accn[g1 * 2 + 0], sq1 = accn[g1 * 2 + 1];
            const float mean1 = sum1 * inv_cnt;
            const float var1 = fmaxf(sq1 * inv_cnt - mean1 * mean1, 0.0f);
            const float rstd1 = rsqrtf(var1 + eps);
            a1 = w[c1] * rstd1;
            b1 = b[c1] - mean1 * a1;
        }
        bs0 = __ldg(cb + c0);
        bs1 = __ldg(cb + c1);
    }
    const float *xb = x + (size_t)n * (size_t)HW * (size_t)C;
    float *outn = out + (size_t)n * (size_t)HW;
    const int quad0 = (blockIdx.x * nwarps + wid) * 4;
    const int step = gridDim.x * nwarps * 4;

    int hw = quad0;
    // Main loop: full quads only, no per-iteration bounds checks.
    for (; hw + 3 < HW; hw += step) {
        float2 v0, v1, v2, v3;
        if (active) {
            const float *p = xb + (size_t)hw * (size_t)C + c0;
            v0 = __ldcs(reinterpret_cast<const float2 *>(p));
            v1 = __ldcs(reinterpret_cast<const float2 *>(p + C));
            v2 = __ldcs(reinterpret_cast<const float2 *>(p + 2 * C));
            v3 = __ldcs(reinterpret_cast<const float2 *>(p + 3 * C));
        } else {
            v0 = v1 = v2 = v3 = make_float2(0.0f, 0.0f);
        }
        float m0, m1, m2, m3, s0, s1, s2, s3;
        quad_lse(v0, v1, v2, v3, a0, b0, a1, b1, bs0, bs1, active,
                 m0, m1, m2, m3, s0, s1, s2, s3);
        // Every lane holds all 4 results after xor butterflies; lane 0 stores
        // one vectorized value (hw % 4 == 0 => 16B aligned).
        if (lane == 0) {
            float4 o4;
            o4.x = m0 + logf(s0);
            o4.y = m1 + logf(s1);
            o4.z = m2 + logf(s2);
            o4.w = m3 + logf(s3);
            __stcs(reinterpret_cast<float4 *>(outn + hw), o4);
        }
    }
    // Tail: masked loads, scalar stores.
    for (; hw < HW; hw += step) {
        float2 v0 = make_float2(0.f, 0.f), v1 = v0, v2 = v0, v3 = v0;
        if (active) {
            const float *p = xb + (size_t)hw * (size_t)C + c0;
            v0 = __ldg(reinterpret_cast<const float2 *>(p));
            if (hw + 1 < HW) v1 = __ldg(reinterpret_cast<const float2 *>(p + C));
            if (hw + 2 < HW) v2 = __ldg(reinterpret_cast<const float2 *>(p + 2 * C));
        }
        float m0, m1, m2, m3, s0, s1, s2, s3;
        quad_lse(v0, v1, v2, v3, a0, b0, a1, b1, bs0, bs1, active,
                 m0, m1, m2, m3, s0, s1, s2, s3);
        if (lane == 0) {
            outn[hw] = m0 + logf(s0);
            if (hw + 1 < HW) outn[hw + 1] = m1 + logf(s1);
            if (hw + 2 < HW) outn[hw + 2] = m2 + logf(s2);
        }
    }
}

__device__ __forceinline__ void lse_update(float r, float &m, float &s) {
    if (r > m) {
        s = s * __expf(m - r) + 1.0f;
        m = r;
    } else {
        s += __expf(r - m);
    }
}

// ---------------------------------------------------------------------------
// Kernel 2 (generic fallback): one warp per spatial position, C <= 32*MAX_CPL.
// ---------------------------------------------------------------------------
__global__ void lse_nhwc_generic_kernel(const float *__restrict__ x,
                                        const float *__restrict__ cb,
                                        const float *__restrict__ w,
                                        const float *__restrict__ b,
                                        const float *__restrict__ acc,
                                        float *__restrict__ out,
                                        int HW, int C, int G, int CperG,
                                        float inv_cnt, float eps, int chunk) {
    const int n = blockIdx.y;
    const int hw0 = blockIdx.x * chunk;
    const int hw1 = min(hw0 + chunk, HW);
    if (hw0 >= hw1) return;
    const int lane = threadIdx.x & 31;
    const int wid = threadIdx.x >> 5;
    const int nwarps = blockDim.x >> 5;
    const int cpl = (C + 31) >> 5;
    float A[MAX_CPL], B[MAX_CPL], CB[MAX_CPL];
    #pragma unroll
    for (int j = 0; j < MAX_CPL; ++j) {
        const int c = lane + (j << 5);
        if (j < cpl && c < C) {
            const int ng = n * G + c / CperG;
            const float mean = acc[ng * 2 + 0] * inv_cnt;
            const float var = fmaxf(acc[ng * 2 + 1] * inv_cnt - mean * mean, 0.0f);
            const float a = w[c] * rsqrtf(var + eps);
            A[j] = a;
            B[j] = b[c] - mean * a;
            CB[j] = __ldg(cb + c);
        }
    }
    const float *xb = x + (size_t)n * (size_t)HW * (size_t)C;
    for (int hw = hw0 + wid; hw < hw1; hw += nwarps) {
        const float *p = xb + (size_t)hw * (size_t)C;
        float m = NEG_INF, s = 0.0f;
        #pragma unroll
        for (int j = 0; j < MAX_CPL; ++j) {
            if (j < cpl) {
                const int c = lane + (j << 5);
                if (c < C) {
                    lse_update(act_res(p[c] + CB[j], A[j], B[j]), m, s);
                }
            }
        }
        #pragma unroll
        for (int o = 16; o > 0; o >>= 1) {
            const float m2 = __shfl_xor_sync(0xffffffffu, m, o);
            const float s2 = __shfl_xor_sync(0xffffffffu, s, o);
            const float M = fmaxf(m, m2);
            float ns = 0.0f;
            if (M != NEG_INF) ns = s * __expf(m - M) + s2 * __expf(m2 - M);
            m = M;
            s = ns;
        }
        if (lane == 0) out[(size_t)n * (size_t)HW + hw] = m + logf(s);
    }
}

torch::Tensor nchw_to_nhwc(torch::Tensor x) {
    TORCH_CHECK(x.is_cuda() && x.dim() == 4 && x.scalar_type() == at::kFloat,
                "x must be CUDA float32 4D");
    auto xc = x.contiguous();
    const int N = xc.size(0), C = xc.size(1), H = xc.size(2), W = xc.size(3);
    const int HW = H * W;
    // Allocate DIRECTLY in channels_last format so the NHWC-ordered buffer
    // needs no further reinterpretation.
    auto out = at::empty({N, C, H, W},
                         xc.options().memory_format(at::MemoryFormat::ChannelsLast));
    TORCH_CHECK(out.is_contiguous(at::MemoryFormat::ChannelsLast),
                "channels_last allocation failed");
    auto stream = at::cuda::getCurrentCUDAStream();
    dim3 grid((HW + 255) / 256, (C + 7) / 8, N);
    nchw_to_nhwc_kernel<<<grid, THREADS, 0, stream>>>(
        xc.data_ptr<float>(), out.data_ptr<float>(), C, HW);
    return out;
}

torch::Tensor fused_tail_nhwc(torch::Tensor x, torch::Tensor cb,
                              torch::Tensor w, torch::Tensor b,
                              int64_t groups, double eps) {
    TORCH_CHECK(x.is_cuda() && x.dim() == 4 && x.scalar_type() == at::kFloat,
                "x must be a CUDA float32 4D tensor");
    TORCH_CHECK(x.is_contiguous(at::MemoryFormat::ChannelsLast),
                "x must be channels_last contiguous");
    const int N = x.size(0), C = x.size(1), H = x.size(2), W = x.size(3);
    const int HW = H * W;
    const int G = (int)groups;
    TORCH_CHECK(C % G == 0, "channels must be divisible by groups");
    TORCH_CHECK(G <= 64, "too many groups");
    TORCH_CHECK(cb.numel() == C && w.numel() == C && b.numel() == C,
                "bad parameter shapes");
    const int CperG = C / G;
    const bool stats_ok = (C % 4 == 0) && (CperG % 4 == 0) && (C / 4) <= THREADS &&
                          (THREADS % (C / 4)) == 0;
    const bool lse_fast_ok = (C % 2 == 0) && (C <= 64);
    TORCH_CHECK(C <= 32 * MAX_CPL, "too many channels");

    auto opts = x.options();
    auto acc = at::zeros({(long)N * G * 2}, opts);
    auto out = at::empty({N, 1, H, W}, opts);
    auto stream = at::cuda::getCurrentCUDAStream();

    const float *xp = x.data_ptr<float>();
    const float *cbp = cb.data_ptr<float>();
    const float *wp = w.data_ptr<float>();
    const float *bp = b.data_ptr<float>();
    float *accp = acc.data_ptr<float>();
    float *outp = out.data_ptr<float>();

    if (stats_ok) {
        const int schunk = 1024;
        dim3 sgrid((HW + schunk - 1) / schunk, N);
        stats_nhwc_kernel<<<sgrid, THREADS, 0, stream>>>(
            reinterpret_cast<const float4 *>(xp),
            reinterpret_cast<const float4 *>(cbp),
            accp, HW, C, G, CperG, schunk);
    } else {
        stats_nhwc_generic_kernel<<<N * G, THREADS, 0, stream>>>(
            xp, cbp, accp, HW, C, G, CperG);
    }

    const float inv_cnt = 1.0f / (float)((long long)CperG * (long long)HW);
    if (lse_fast_ok) {
        // fast path: warp per 4 spatial positions, two-phase reduction
        const int nwarps = THREADS / 32;
        const int iters = 8;
        dim3 lgrid((HW + 4 * nwarps * iters - 1) / (4 * nwarps * iters), N);
        lse_nhwc_kernel<<<lgrid, THREADS, 0, stream>>>(
            xp, cbp, wp, bp, accp, outp, HW, C, G, CperG, inv_cnt, (float)eps);
    } else {
        const int lchunk = 512;
        dim3 lgrid((HW + lchunk - 1) / lchunk, N);
        lse_nhwc_generic_kernel<<<lgrid, THREADS, 0, stream>>>(
            xp, cbp, wp, bp, accp, outp, HW, C, G, CperG, inv_cnt, (float)eps,
            lchunk);
    }
    return out;
}
"""

_cpp_src = (
    "torch::Tensor nchw_to_nhwc(torch::Tensor x);\n"
    "torch::Tensor fused_tail_nhwc(torch::Tensor x, torch::Tensor cb, torch::Tensor w, torch::Tensor b, int64_t groups, double eps);"
)

_fused = load_inline(
    name="fused_tail_nhwc_v6",
    cpp_sources=_cpp_src,
    cuda_sources=_cuda_src,
    functions=["nchw_to_nhwc", "fused_tail_nhwc"],
    verbose=False,
    extra_cflags=["-O3"],
    extra_cuda_cflags=["-O3"],
)


class ModelNew(nn.Module):
    """
    Optimized model: custom NCHW->NHWC input transpose, channels_last cuDNN
    convolution WITHOUT bias (the conv bias is fused into the custom tail
    kernels, eliminating cuDNN's separate full-tensor bias-add kernel), then
    two custom CUDA kernels that fuse GroupNorm statistics +
    GroupNorm/Tanh/HardSwish/Residual/LogSumExp.
    """
    def __init__(self, in_channels, out_channels, kernel_size, groups, eps=1e-5):
        super(ModelNew, self).__init__()
        self.conv = nn.Conv2d(in_channels, out_channels, kernel_size)
        self.group_norm = nn.GroupNorm(groups, out_channels, eps=eps)
        self.tanh = nn.Tanh()
        self.hard_swish = nn.Hardswish()
        self.groups = groups
        self.eps = float(eps)
        C = out_channels
        g = groups
        # Static check that the custom NHWC tail supports this configuration.
        self._tail_ok = (
            g > 0
            and C % g == 0
            and C <= 128
            and g <= 64
        )

    def forward(self, x):
        if (
            x.is_cuda
            and x.dtype == torch.float32
            and self._tail_ok
            and self.conv.bias is not None
        ):
            # Custom fast transpose to channels_last (much faster than torch's
            # .contiguous(memory_format=channels_last)).
            x_cl = _fused.nchw_to_nhwc(x)
            # NOTE: do NOT cache this - weights may be updated between calls.
            # It is only a tiny (out_ch, in_ch, k, k) copy.
            w_cl = self.conv.weight.contiguous(memory_format=torch.channels_last)
            # bias=None: the conv bias is applied inside the fused tail kernels.
            x_conv = F.conv2d(
                x_cl,
                w_cl,
                None,
                self.conv.stride,
                self.conv.padding,
                self.conv.dilation,
                self.conv.groups,
            )
            if x_conv.is_contiguous(memory_format=torch.channels_last):
                return _fused.fused_tail_nhwc(
                    x_conv,
                    self.conv.bias,
                    self.group_norm.weight,
                    self.group_norm.bias,
                    self.groups,
                    self.eps,
                )
        # Reference fallback (CPU tensors or unsupported configurations).
        x_conv = self.conv(x)
        x_norm = self.group_norm(x_conv)
        x_tanh = torch.tanh(x_norm)
        x_hard_swish = F.hardswish(x_tanh)
        x_res = x_conv + x_hard_swish
        return torch.logsumexp(x_res, dim=1, keepdim=True)
```