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


_cpp_source = r"""
#include <torch/extension.h>

torch::Tensor channels_last_input(torch::Tensor x);
torch::Tensor fused_conv_mish(
    torch::Tensor x, torch::Tensor weight, torch::Tensor bias,
    double s1, double s2);
torch::Tensor fused_bias_mish(
    torch::Tensor x, torch::Tensor bias, double s1, double s2);
"""

_cuda_source = r"""
#include <torch/extension.h>
#include <ATen/cuda/CUDAContext.h>
#include <c10/cuda/CUDAGuard.h>
#include <c10/cuda/CUDAException.h>
#include <cuda_runtime.h>
#include <mma.h>
#include <algorithm>
#include <cstdint>

namespace wmma = nvcuda::wmma;

__device__ __forceinline__ float round_tf32(float value) {
#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 800
    unsigned int bits;
    asm("cvt.rn.tf32.f32 %0, %1;" : "=r"(bits) : "f"(value));
    return __uint_as_float(bits);
#else
    return value;
#endif
}

__global__ void nchw8_to_nhwc(
    const float* __restrict__ src,
    float* __restrict__ dst,
    int64_t spatial,
    int64_t tiles_per_image
) {
    __shared__ float tile[8][516];

    const int tid = threadIdx.x;
    const int channel = tid >> 5;
    const int lane = tid & 31;
    const int64_t batch = blockIdx.x / tiles_per_image;
    const int64_t first_pixel = (blockIdx.x % tiles_per_image) * 512;

    #pragma unroll
    for (int j = 0; j < 4; ++j) {
        const int p = lane * 4 + j * 128;
        const float4 v = *reinterpret_cast<const float4*>(
            src + (batch * 8 + channel) * spatial + first_pixel + p);

        tile[channel][p] = v.x;
        tile[channel][p + 1] = v.y;
        tile[channel][p + 2] = v.z;
        tile[channel][p + 3] = v.w;
    }

    __syncthreads();

    const int c = (tid & 1) * 4;

    #pragma unroll
    for (int j = 0; j < 4; ++j) {
        const int p = (tid >> 1) + j * 128;
        const float4 v = make_float4(
            tile[c][p], tile[c + 1][p],
            tile[c + 2][p], tile[c + 3][p]);

        *reinterpret_cast<float4*>(
            dst + (batch * spatial + first_pixel + p) * 8 + c) = v;
    }
}

__device__ __forceinline__ float activate(
    float x, float bias, float s1, float s2
) {
    x = x + bias;
    x = x - s1;
    x = x - s2;

    if (x > 20.0f)
        return x;

    const float e = __expf(x);
    const float q = e * (e + 2.0f);
    return x * __fdividef(q, q + 2.0f);
}

__device__ __forceinline__ float4 activate4(
    float4 x, float4 b, float s1, float s2
) {
    x.x = activate(x.x, b.x, s1, s2);
    x.y = activate(x.y, b.y, s1, s2);
    x.z = activate(x.z, b.z, s1, s2);
    x.w = activate(x.w, b.w, s1, s2);
    return x;
}

__global__ void pack_weights(
    const float* __restrict__ weight,
    float* __restrict__ packed
) {
    const int i = blockIdx.x * blockDim.x + threadIdx.x;
    if (i < 64 * 80) {
        const int c = i / 80;
        const int k = i % 80;
        const float v = k < 72 ? weight[c * 72 + k] : 0.0f;
        const float hi = round_tf32(v);
        packed[i] = hi;
        packed[64 * 80 + i] = round_tf32(v - hi);
    }
}

__global__ __launch_bounds__(256) void conv8_64_mish(
    const float* __restrict__ input,
    const float* __restrict__ packed_weight,
    const float4* __restrict__ bias,
    float4* __restrict__ output,
    float s1,
    float s2
) {
#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 800
    constexpr int M = 64;
    constexpr int LD = 80;
    constexpr int Spatial = 254 * 254;

    __shared__ __align__(32) float storage[2 * M * LD];
    float* const ah = storage;
    float* const al = storage + M * LD;

    const int tid = threadIdx.x;
    const int first = blockIdx.x * M;
    const int batch = blockIdx.y;

    for (int i = tid; i < M * 18; i += 256) {
        const int row = i / 18;
        const int kv = i - row * 18;
        const int pixel = first + row;
        float4 v = make_float4(0.f, 0.f, 0.f, 0.f);

        if (pixel < Spatial) {
            const int oh = pixel / 254;
            const int ow = pixel - oh * 254;
            const int position = kv >> 1;
            const int kh = position / 3;
            const int kw = position - kh * 3;
            const int c = (kv & 1) * 4;

            const int64_t offset =
                ((int64_t(batch) * 256 + oh + kh) * 256 + ow + kw) * 8 + c;
            v = *reinterpret_cast<const float4*>(input + offset);
        }

        const float4 hi = make_float4(
            round_tf32(v.x), round_tf32(v.y),
            round_tf32(v.z), round_tf32(v.w));
        const float4 lo = make_float4(
            round_tf32(v.x - hi.x), round_tf32(v.y - hi.y),
            round_tf32(v.z - hi.z), round_tf32(v.w - hi.w));

        *reinterpret_cast<float4*>(ah + row * LD + kv * 4) = hi;
        *reinterpret_cast<float4*>(al + row * LD + kv * 4) = lo;
    }

    __syncthreads();

    const int warp = tid >> 5;
    const int row = (warp >> 1) * 16;
    const int col = (warp & 1) * 32;

    wmma::fragment<wmma::accumulator, 16, 16, 8, float> acc0, acc1;
    wmma::fill_fragment(acc0, 0.0f);
    wmma::fill_fragment(acc1, 0.0f);

    const float* const bh = packed_weight;
    const float* const bl = packed_weight + 64 * LD;

    #pragma unroll
    for (int k = 0; k < 72; k += 8) {
        wmma::fragment<
            wmma::matrix_a, 16, 16, 8,
            wmma::precision::tf32, wmma::row_major> a_hi, a_lo;
        wmma::fragment<
            wmma::matrix_b, 16, 16, 8,
            wmma::precision::tf32, wmma::col_major> b_hi, b_lo;

        wmma::load_matrix_sync(a_hi, ah + row * LD + k, LD);
        wmma::load_matrix_sync(a_lo, al + row * LD + k, LD);

        wmma::load_matrix_sync(b_hi, bh + col * LD + k, LD);
        wmma::load_matrix_sync(b_lo, bl + col * LD + k, LD);
        wmma::mma_sync(acc0, a_hi, b_hi, acc0);
        wmma::mma_sync(acc0, a_lo, b_hi, acc0);
        wmma::mma_sync(acc0, a_hi, b_lo, acc0);

        wmma::load_matrix_sync(b_hi, bh + (col + 16) * LD + k, LD);
        wmma::load_matrix_sync(b_lo, bl + (col + 16) * LD + k, LD);
        wmma::mma_sync(acc1, a_hi, b_hi, acc1);
        wmma::mma_sync(acc1, a_lo, b_hi, acc1);
        wmma::mma_sync(acc1, a_hi, b_lo, acc1);
    }

    __syncthreads();

    wmma::store_matrix_sync(
        storage + row * 64 + col, acc0, 64, wmma::mem_row_major);
    wmma::store_matrix_sync(
        storage + row * 64 + col + 16, acc1, 64, wmma::mem_row_major);

    __syncthreads();

    const float4 b = bias[tid & 15];

    #pragma unroll
    for (int j = 0; j < 4; ++j) {
        const int i = tid + j * 256;
        const int pixel = first + (i >> 4);
        if (pixel < Spatial) {
            float4 v = reinterpret_cast<const float4*>(storage)[i];
            v = activate4(v, b, s1, s2);
            output[(int64_t(batch) * Spatial + pixel) * 16 + (i & 15)] = v;
        }
    }
#endif
}

__global__ void bias_mish64(
    float4* __restrict__ x,
    const float4* __restrict__ bias,
    int64_t vectors,
    float s1,
    float s2
) {
    const float4 b = bias[threadIdx.x & 15];
    const int64_t step = int64_t(gridDim.x) * 1024;

    for (int64_t base = int64_t(blockIdx.x) * 1024;
         base < vectors; base += step) {
        const int64_t i = base + threadIdx.x;

        if (base + 1024 <= vectors) {
            float4 a0 = x[i];
            float4 a1 = x[i + 256];
            float4 a2 = x[i + 512];
            float4 a3 = x[i + 768];

            a0 = activate4(a0, b, s1, s2);
            a1 = activate4(a1, b, s1, s2);
            a2 = activate4(a2, b, s1, s2);
            a3 = activate4(a3, b, s1, s2);

            x[i] = a0;
            x[i + 256] = a1;
            x[i + 512] = a2;
            x[i + 768] = a3;
        } else {
            #pragma unroll
            for (int j = 0; j < 4; ++j) {
                const int64_t k = i + j * 256;
                if (k < vectors)
                    x[k] = activate4(x[k], b, s1, s2);
            }
        }
    }
}

__global__ void bias_mish_vector(
    float4* __restrict__ x,
    const float4* __restrict__ bias,
    int64_t vectors,
    int channel_vectors,
    float s1,
    float s2
) {
    for (int64_t i = int64_t(blockIdx.x) * blockDim.x + threadIdx.x;
         i < vectors; i += int64_t(blockDim.x) * gridDim.x) {
        x[i] = activate4(x[i], bias[i % channel_vectors], s1, s2);
    }
}

__global__ void bias_mish_scalar(
    float* __restrict__ x,
    const float* __restrict__ bias,
    int64_t count,
    int channels,
    float s1,
    float s2
) {
    for (int64_t i = int64_t(blockIdx.x) * blockDim.x + threadIdx.x;
         i < count; i += int64_t(blockDim.x) * gridDim.x) {
        x[i] = activate(x[i], bias[i % channels], s1, s2);
    }
}

torch::Tensor channels_last_input(torch::Tensor x) {
    TORCH_CHECK(x.is_cuda() && x.dim() == 4, "Expected CUDA NCHW tensor");

    if (x.is_contiguous(at::MemoryFormat::ChannelsLast))
        return x;

    const c10::cuda::CUDAGuard guard(x.device());
    const int64_t spatial = x.size(2) * x.size(3);

    if (x.scalar_type() != at::kFloat ||
        x.size(1) != 8 ||
        !x.is_contiguous() ||
        spatial == 0 ||
        spatial % 512 != 0 ||
        reinterpret_cast<uintptr_t>(x.data_ptr()) % 16 != 0) {
        return x.contiguous(at::MemoryFormat::ChannelsLast);
    }

    auto output = torch::empty(
        x.sizes(), x.options().memory_format(at::MemoryFormat::ChannelsLast));

    const int64_t tiles = spatial / 512;
    const int64_t blocks = x.size(0) * tiles;

    if (blocks != 0) {
        nchw8_to_nhwc<<<blocks, 256, 0, at::cuda::getCurrentCUDAStream()>>>(
            x.data_ptr<float>(), output.data_ptr<float>(), spatial, tiles);
        C10_CUDA_KERNEL_LAUNCH_CHECK();
    }

    return output;
}

torch::Tensor fused_conv_mish(
    torch::Tensor x, torch::Tensor weight, torch::Tensor bias,
    double s1, double s2
) {
    TORCH_CHECK(
        x.is_cuda() && weight.is_cuda() && bias.is_cuda() &&
        x.device() == weight.device() && x.device() == bias.device(),
        "Expected tensors on the same CUDA device");
    TORCH_CHECK(
        x.scalar_type() == at::kFloat &&
        weight.scalar_type() == at::kFloat &&
        bias.scalar_type() == at::kFloat,
        "Expected float32 tensors");
    TORCH_CHECK(
        x.dim() == 4 && x.size(1) == 8 &&
        x.size(2) == 256 && x.size(3) == 256 &&
        x.size(0) <= 65535 &&
        x.is_contiguous(at::MemoryFormat::ChannelsLast),
        "Unsupported input");
    TORCH_CHECK(
        weight.dim() == 4 && weight.size(0) == 64 &&
        weight.size(1) == 8 && weight.size(2) == 3 &&
        weight.size(3) == 3 &&
        weight.is_contiguous(at::MemoryFormat::ChannelsLast) &&
        bias.is_contiguous() && bias.numel() == 64,
        "Unsupported weights");

    const c10::cuda::CUDAGuard guard(x.device());
    TORCH_CHECK(
        at::cuda::getCurrentDeviceProperties()->major >= 8,
        "TF32 convolution requires compute capability 8.0 or newer");

    auto output = torch::empty(
        {x.size(0), 64, 254, 254},
        x.options().memory_format(at::MemoryFormat::ChannelsLast));

    if (x.size(0) == 0)
        return output;

    auto packed = torch::empty({2, 64, 80}, weight.options());
    const auto stream = at::cuda::getCurrentCUDAStream();

    pack_weights<<<20, 256, 0, stream>>>(
        weight.data_ptr<float>(), packed.data_ptr<float>());
    C10_CUDA_KERNEL_LAUNCH_CHECK();

    const dim3 grid((254 * 254 + 63) / 64, x.size(0));
    conv8_64_mish<<<grid, 256, 0, stream>>>(
        x.data_ptr<float>(),
        packed.data_ptr<float>(),
        reinterpret_cast<const float4*>(bias.data_ptr<float>()),
        reinterpret_cast<float4*>(output.data_ptr<float>()),
        static_cast<float>(s1), static_cast<float>(s2));
    C10_CUDA_KERNEL_LAUNCH_CHECK();

    return output;
}

torch::Tensor fused_bias_mish(
    torch::Tensor x, torch::Tensor bias, double s1, double s2
) {
    TORCH_CHECK(
        x.is_cuda() && bias.is_cuda() && x.device() == bias.device(),
        "Expected tensors on the same CUDA device");
    TORCH_CHECK(
        x.scalar_type() == at::kFloat &&
        bias.scalar_type() == at::kFloat &&
        x.dim() == 4 &&
        x.is_contiguous(at::MemoryFormat::ChannelsLast) &&
        bias.is_contiguous() && bias.numel() == x.size(1),
        "Invalid tensors");

    const c10::cuda::CUDAGuard guard(x.device());
    const int64_t count = x.numel();
    if (count == 0)
        return x;

    const int channels = static_cast<int>(x.size(1));
    const auto stream = at::cuda::getCurrentCUDAStream();
    const bool aligned =
        reinterpret_cast<uintptr_t>(x.data_ptr()) % 16 == 0 &&
        reinterpret_cast<uintptr_t>(bias.data_ptr()) % 16 == 0;

    if (channels == 64 && aligned) {
        const int64_t vectors = count / 4;
        const int blocks = static_cast<int>(
            std::min<int64_t>((vectors + 1023) / 1024, 2048));

        bias_mish64<<<blocks, 256, 0, stream>>>(
            reinterpret_cast<float4*>(x.data_ptr<float>()),
            reinterpret_cast<const float4*>(bias.data_ptr<float>()),
            vectors, static_cast<float>(s1), static_cast<float>(s2));
    } else if (channels % 4 == 0 && aligned) {
        const int64_t vectors = count / 4;
        const int blocks = static_cast<int>(
            std::min<int64_t>((vectors + 255) / 256, 2048));

        bias_mish_vector<<<blocks, 256, 0, stream>>>(
            reinterpret_cast<float4*>(x.data_ptr<float>()),
            reinterpret_cast<const float4*>(bias.data_ptr<float>()),
            vectors, channels / 4,
            static_cast<float>(s1), static_cast<float>(s2));
    } else {
        const int blocks = static_cast<int>(
            std::min<int64_t>((count + 255) / 256, 2048));

        bias_mish_scalar<<<blocks, 256, 0, stream>>>(
            x.data_ptr<float>(), bias.data_ptr<float>(), count, channels,
            static_cast<float>(s1), static_cast<float>(s2));
    }

    C10_CUDA_KERNEL_LAUNCH_CHECK();
    return x;
}
"""

_ops = load_inline(
    name="fused_conv8_64_mish_tf32x3_v6",
    cpp_sources=_cpp_source,
    cuda_sources=_cuda_source,
    functions=[
        "channels_last_input",
        "fused_conv_mish",
        "fused_bias_mish",
    ],
    extra_cflags=["-O3"],
    extra_cuda_cflags=["-O3", "--fmad=false"],
    verbose=False,
)


class ModelNew(nn.Module):
    def __init__(
        self,
        in_channels,
        out_channels,
        kernel_size,
        subtract_value_1,
        subtract_value_2,
    ):
        super().__init__()
        self.conv = nn.Conv2d(in_channels, out_channels, kernel_size)
        self.conv.to(memory_format=torch.channels_last)
        self.subtract_value_1 = subtract_value_1
        self.subtract_value_2 = subtract_value_2

    def forward(self, x):
        if (
            not x.is_cuda
            or x.dtype != torch.float32
            or self.conv.weight.dtype != torch.float32
            or self.conv.bias is None
            or self.conv.bias.dtype != torch.float32
            or torch.is_autocast_enabled()
        ):
            return F.mish(
                (self.conv(x) - self.subtract_value_1)
                - self.subtract_value_2
            )

        x = _ops.channels_last_input(x)
        weight = self.conv.weight.contiguous(
            memory_format=torch.channels_last
        )
        bias = self.conv.bias.contiguous()

        if (
            x.shape[1:] == (8, 256, 256)
            and x.shape[0] <= 65535
            and weight.shape == (64, 8, 3, 3)
            and self.conv.stride == (1, 1)
            and self.conv.padding == (0, 0)
            and self.conv.dilation == (1, 1)
            and self.conv.groups == 1
            and torch.cuda.get_device_capability(x.device)[0] >= 8
        ):
            return _ops.fused_conv_mish(
                x,
                weight,
                bias,
                self.subtract_value_1,
                self.subtract_value_2,
            )

        y = F.conv2d(
            x,
            weight,
            None,
            self.conv.stride,
            self.conv.padding,
            self.conv.dilation,
            self.conv.groups,
        ).contiguous(memory_format=torch.channels_last)

        return _ops.fused_bias_mish(
            y,
            bias,
            self.subtract_value_1,
            self.subtract_value_2,
        )
```