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


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

torch::Tensor depthwise3x3_cuda(
    torch::Tensor input,
    torch::Tensor weight,
    c10::optional<torch::Tensor> bias);
"""

_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 <cstdint>

template<bool HasBias, int Rows>
__global__ __launch_bounds__(128)
void depthwise3x3_aligned_kernel(
    const float* __restrict__ input,
    const float* __restrict__ weight,
    const float* __restrict__ bias,
    float* __restrict__ output,
    int channels,
    int height,
    int width,
    int output_height,
    int64_t planes)
{
    const int x = (blockIdx.x * 128 + threadIdx.x) * 4;
    const int y = blockIdx.y * Rows;
    if (x >= width)
        return;

    const bool has_halo = x + 4 < width;

    for (int64_t plane = blockIdx.z; plane < planes; plane += gridDim.z) {
        const int channel = plane % channels;
        const float* w = weight + channel * 9;

        const float w00 = w[0], w01 = w[1], w02 = w[2];
        const float w10 = w[3], w11 = w[4], w12 = w[5];
        const float w20 = w[6], w21 = w[7], w22 = w[8];

        float b = 0.0f;
        if (HasBias)
            b = bias[channel];

        const float* src =
            input + plane * int64_t(height) * width + int64_t(y) * width + x;
        float* dst =
            output + plane * int64_t(output_height) * width
            + int64_t(y) * width + x;

        float4 r0 = *reinterpret_cast<const float4*>(src);
        float4 r1 = *reinterpret_cast<const float4*>(src + width);

        float2 t0 = make_float2(0.0f, 0.0f);
        float2 t1 = make_float2(0.0f, 0.0f);
        if (has_halo) {
            t0 = *reinterpret_cast<const float2*>(src + 4);
            t1 = *reinterpret_cast<const float2*>(src + width + 4);
        }

        #pragma unroll
        for (int row = 0; row < Rows; ++row) {
            if (y + row < output_height) {
                const float* next = src + int64_t(row + 2) * width;
                const float4 r2 = *reinterpret_cast<const float4*>(next);

                float2 t2 = make_float2(0.0f, 0.0f);
                if (has_halo)
                    t2 = *reinterpret_cast<const float2*>(next + 4);

                float a0 = r0.x * w00;
                float a1 = r0.y * w00;
                float a2 = r0.z * w00;
                float a3 = r0.w * w00;

                a0 = fmaf(r0.y, w01, a0);
                a1 = fmaf(r0.z, w01, a1);
                a2 = fmaf(r0.w, w01, a2);
                a3 = fmaf(t0.x, w01, a3);

                a0 = fmaf(r0.z, w02, a0);
                a1 = fmaf(r0.w, w02, a1);
                a2 = fmaf(t0.x, w02, a2);
                a3 = fmaf(t0.y, w02, a3);

                a0 = fmaf(r1.x, w10, a0);
                a1 = fmaf(r1.y, w10, a1);
                a2 = fmaf(r1.z, w10, a2);
                a3 = fmaf(r1.w, w10, a3);

                a0 = fmaf(r1.y, w11, a0);
                a1 = fmaf(r1.z, w11, a1);
                a2 = fmaf(r1.w, w11, a2);
                a3 = fmaf(t1.x, w11, a3);

                a0 = fmaf(r1.z, w12, a0);
                a1 = fmaf(r1.w, w12, a1);
                a2 = fmaf(t1.x, w12, a2);
                a3 = fmaf(t1.y, w12, a3);

                a0 = fmaf(r2.x, w20, a0);
                a1 = fmaf(r2.y, w20, a1);
                a2 = fmaf(r2.z, w20, a2);
                a3 = fmaf(r2.w, w20, a3);

                a0 = fmaf(r2.y, w21, a0);
                a1 = fmaf(r2.z, w21, a1);
                a2 = fmaf(r2.w, w21, a2);
                a3 = fmaf(t2.x, w21, a3);

                a0 = fmaf(r2.z, w22, a0);
                a1 = fmaf(r2.w, w22, a1);
                a2 = fmaf(t2.x, w22, a2);
                a3 = fmaf(t2.y, w22, a3);

                if (HasBias) {
                    a0 += b;
                    a1 += b;
                    a2 += b;
                    a3 += b;
                }

                *reinterpret_cast<float4*>(dst + int64_t(row) * width) =
                    make_float4(a0, a1, a2, a3);

                r0 = r1;
                r1 = r2;
                t0 = t1;
                t1 = t2;
            }
        }
    }
}

torch::Tensor depthwise3x3_cuda(
    torch::Tensor input,
    torch::Tensor weight,
    c10::optional<torch::Tensor> bias)
{
    const c10::cuda::CUDAGuard device_guard(input.device());

    const int64_t batch = input.size(0);
    const int channels = input.size(1);
    const int height = input.size(2);
    const int width = input.size(3);
    const int output_height = height - 2;
    const int64_t planes = batch * channels;

    // Padding the physical row stride keeps every vector store aligned.
    auto storage = torch::empty(
        {batch, channels, output_height, width}, input.options());
    auto output = storage.narrow(3, 0, width - 2);

    if (output.numel() == 0)
        return output;

    constexpr int rows = 8;
    const dim3 grid(
        (width + 511) / 512,
        (output_height + rows - 1) / rows,
        static_cast<unsigned int>(planes < 65535 ? planes : 65535));

    const auto stream = at::cuda::getCurrentCUDAStream();

    if (bias.has_value()) {
        depthwise3x3_aligned_kernel<true, rows><<<grid, 128, 0, stream>>>(
            input.data_ptr<float>(),
            weight.data_ptr<float>(),
            bias.value().data_ptr<float>(),
            storage.data_ptr<float>(),
            channels, height, width, output_height, planes);
    } else {
        depthwise3x3_aligned_kernel<false, rows><<<grid, 128, 0, stream>>>(
            input.data_ptr<float>(),
            weight.data_ptr<float>(),
            nullptr,
            storage.data_ptr<float>(),
            channels, height, width, output_height, planes);
    }

    C10_CUDA_KERNEL_LAUNCH_CHECK();
    return output;
}
"""

_depthwise_extension = load_inline(
    name="depthwise3x3_aligned_vec4_rows8_v3",
    cpp_sources=_cpp_source,
    cuda_sources=_cuda_source,
    functions=["depthwise3x3_cuda"],
    extra_cflags=["-O3"],
    extra_cuda_cflags=["-O3"],
    verbose=False,
)


class ModelNew(nn.Module):
    def __init__(
        self,
        in_channels: int,
        out_channels: int,
        kernel_size: int,
        stride: int = 1,
        padding: int = 0,
        bias: bool = False,
    ):
        super().__init__()
        self.conv2d = nn.Conv2d(
            in_channels,
            out_channels,
            kernel_size=(kernel_size, kernel_size),
            stride=stride,
            padding=padding,
            groups=in_channels,
            bias=bias,
        )

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        conv = self.conv2d
        if (
            x.is_cuda
            and x.dtype == torch.float32
            and x.ndim == 4
            and x.is_contiguous()
            and x.size(1) == conv.in_channels
            and x.size(2) >= 3
            and x.size(3) >= 4
            and x.size(3) % 4 == 0
            and (x.size(2) + 5) // 8 <= 65535
            and x.data_ptr() % 16 == 0
            and conv.in_channels == conv.out_channels
            and conv.kernel_size == (3, 3)
            and conv.stride == (1, 1)
            and conv.padding == (0, 0)
            and conv.dilation == (1, 1)
            and conv.weight.dtype == torch.float32
            and conv.weight.device == x.device
            and conv.weight.is_contiguous()
            and (
                conv.bias is None
                or (
                    conv.bias.dtype == torch.float32
                    and conv.bias.device == x.device
                    and conv.bias.is_contiguous()
                )
            )
            and not torch.is_autocast_enabled()
        ):
            return _depthwise_extension.depthwise3x3_cuda(
                x, conv.weight, conv.bias
            )
        return conv(x)
```