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


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

torch::Tensor depthwise_cuda(
    torch::Tensor x,
    torch::Tensor weight,
    c10::optional<torch::Tensor> bias,
    int64_t stride,
    int64_t padding);
"""

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

// Each thread computes a 4x4 output tile, reusing input values in registers.
__global__ void depthwise_3x3_kernel(
    const float* __restrict__ x,
    const float* __restrict__ weight,
    const float* __restrict__ bias,
    float* __restrict__ out,
    int channels,
    int out_channels,
    int height,
    int width,
    int out_height,
    int out_width)
{
    const int ox = (blockIdx.x * blockDim.x + threadIdx.x) * 4;
    const int oy = (blockIdx.y * blockDim.y + threadIdx.y) * 4;
    if (ox >= out_width || oy >= out_height) return;

    const int nc = blockIdx.z;
    const int oc = nc % out_channels;
    const int n = nc / out_channels;
    const int ic = oc / (out_channels / channels);

    const float* src =
        x + ((int64_t)n * channels + ic) * height * width;
    float* dst = out + (int64_t)nc * out_height * out_width;

    float w[9];
    #pragma unroll
    for (int k = 0; k < 9; ++k)
        w[k] = weight[oc * 9 + k];

    const float b = bias ? bias[oc] : 0.0f;

    if (ox + 3 < out_width && oy + 3 < out_height) {
        float acc[4][4];

        #pragma unroll
        for (int i = 0; i < 4; ++i) {
            #pragma unroll
            for (int j = 0; j < 4; ++j)
                acc[i][j] = 0.0f;
        }

        #pragma unroll
        for (int r = 0; r < 6; ++r) {
            const float* row = src + (int64_t)(oy + r) * width + ox;
            const float2 p0 = *reinterpret_cast<const float2*>(row);
            const float2 p1 = *reinterpret_cast<const float2*>(row + 2);
            const float2 p2 = *reinterpret_cast<const float2*>(row + 4);

            #pragma unroll
            for (int i = 0; i < 4; ++i) {
                if (r >= i && r < i + 3) {
                    const int k = (r - i) * 3;
                    acc[i][0] = fmaf(p0.x, w[k],     acc[i][0]);
                    acc[i][0] = fmaf(p0.y, w[k + 1], acc[i][0]);
                    acc[i][0] = fmaf(p1.x, w[k + 2], acc[i][0]);

                    acc[i][1] = fmaf(p0.y, w[k],     acc[i][1]);
                    acc[i][1] = fmaf(p1.x, w[k + 1], acc[i][1]);
                    acc[i][1] = fmaf(p1.y, w[k + 2], acc[i][1]);

                    acc[i][2] = fmaf(p1.x, w[k],     acc[i][2]);
                    acc[i][2] = fmaf(p1.y, w[k + 1], acc[i][2]);
                    acc[i][2] = fmaf(p2.x, w[k + 2], acc[i][2]);

                    acc[i][3] = fmaf(p1.y, w[k],     acc[i][3]);
                    acc[i][3] = fmaf(p2.x, w[k + 1], acc[i][3]);
                    acc[i][3] = fmaf(p2.y, w[k + 2], acc[i][3]);
                }
            }
        }

        #pragma unroll
        for (int i = 0; i < 4; ++i) {
            float* row = dst + (int64_t)(oy + i) * out_width + ox;
            *reinterpret_cast<float2*>(row) =
                make_float2(acc[i][0] + b, acc[i][1] + b);
            *reinterpret_cast<float2*>(row + 2) =
                make_float2(acc[i][2] + b, acc[i][3] + b);
        }
    } else {
        #pragma unroll
        for (int i = 0; i < 4; ++i) {
            #pragma unroll
            for (int j = 0; j < 4; ++j) {
                if (oy + i < out_height && ox + j < out_width) {
                    float acc = 0.0f;
                    #pragma unroll
                    for (int ky = 0; ky < 3; ++ky) {
                        #pragma unroll
                        for (int kx = 0; kx < 3; ++kx) {
                            acc = fmaf(
                                src[(int64_t)(oy + i + ky) * width
                                    + ox + j + kx],
                                w[ky * 3 + kx], acc);
                        }
                    }
                    dst[(int64_t)(oy + i) * out_width + ox + j] = acc + b;
                }
            }
        }
    }
}

__global__ void depthwise_generic_kernel(
    const float* __restrict__ x,
    const float* __restrict__ weight,
    const float* __restrict__ bias,
    float* __restrict__ out,
    int64_t count,
    int channels,
    int out_channels,
    int height,
    int width,
    int out_height,
    int out_width,
    int kernel,
    int stride,
    int padding)
{
    for (int64_t index = (int64_t)blockIdx.x * blockDim.x + threadIdx.x;
         index < count;
         index += (int64_t)blockDim.x * gridDim.x) {
        int64_t t = index;
        const int ox = t % out_width;
        t /= out_width;
        const int oy = t % out_height;
        t /= out_height;
        const int oc = t % out_channels;
        const int n = t / out_channels;
        const int ic = oc / (out_channels / channels);

        const int iy = oy * stride - padding;
        const int ix = ox * stride - padding;
        const float* src =
            x + ((int64_t)n * channels + ic) * height * width;
        const float* filter = weight + (int64_t)oc * kernel * kernel;

        float acc = 0.0f;
        for (int ky = 0; ky < kernel; ++ky) {
            const int y = iy + ky;
            if ((unsigned)y >= (unsigned)height) continue;
            for (int kx = 0; kx < kernel; ++kx) {
                const int xx = ix + kx;
                if ((unsigned)xx < (unsigned)width)
                    acc = fmaf(src[(int64_t)y * width + xx],
                               filter[ky * kernel + kx], acc);
            }
        }
        out[index] = acc + (bias ? bias[oc] : 0.0f);
    }
}

torch::Tensor depthwise_cuda(
    torch::Tensor x,
    torch::Tensor weight,
    c10::optional<torch::Tensor> bias,
    int64_t stride,
    int64_t padding)
{
    TORCH_CHECK(x.is_cuda() && weight.is_cuda(), "Expected CUDA tensors");
    TORCH_CHECK(x.scalar_type() == at::kFloat &&
                weight.scalar_type() == at::kFloat, "Expected float32 tensors");
    TORCH_CHECK(x.is_contiguous() && weight.is_contiguous(),
                "Expected contiguous tensors");
    TORCH_CHECK(x.dim() == 4 && weight.dim() == 4,
                "Expected four-dimensional tensors");
    TORCH_CHECK(x.device() == weight.device(), "Device mismatch");
    TORCH_CHECK(stride > 0 && padding >= 0, "Invalid convolution parameters");

    const c10::cuda::CUDAGuard device_guard(x.device());

    const int64_t n = x.size(0);
    const int c = x.size(1);
    const int h = x.size(2);
    const int w = x.size(3);
    const int oc = weight.size(0);
    const int k = weight.size(2);

    TORCH_CHECK(c > 0 && oc > 0 && oc % c == 0 &&
                weight.size(1) == 1 && weight.size(3) == k,
                "Invalid depthwise convolution weights");
    TORCH_CHECK(h + 2 * padding >= k && w + 2 * padding >= k,
                "Kernel exceeds padded input");

    const int oh = (h + 2 * padding - k) / stride + 1;
    const int ow = (w + 2 * padding - k) / stride + 1;

    const float* bias_ptr = nullptr;
    if (bias.has_value()) {
        TORCH_CHECK(bias->device() == x.device() &&
                    bias->scalar_type() == at::kFloat &&
                    bias->is_contiguous() && bias->numel() == oc,
                    "Invalid bias");
        bias_ptr = bias->data_ptr<float>();
    }

    auto out = torch::empty({n, oc, oh, ow}, x.options());
    if (out.numel() == 0) return out;

    cudaStream_t stream = at::cuda::getCurrentCUDAStream();

    if (k == 3 && stride == 1 && padding == 0 &&
        (w % 2 == 0) && n * oc <= 65535) {
        const dim3 block(32, 4);
        const dim3 grid((ow + 127) / 128, (oh + 15) / 16, n * oc);
        depthwise_3x3_kernel<<<grid, block, 0, stream>>>(
            x.data_ptr<float>(), weight.data_ptr<float>(), bias_ptr,
            out.data_ptr<float>(), c, oc, h, w, oh, ow);
    } else {
        const int threads = 256;
        const int blocks = static_cast<int>(
            std::min<int64_t>((out.numel() + threads - 1) / threads, 65535));
        depthwise_generic_kernel<<<blocks, threads, 0, stream>>>(
            x.data_ptr<float>(), weight.data_ptr<float>(), bias_ptr,
            out.data_ptr<float>(), out.numel(), c, oc, h, w, oh, ow,
            k, stride, padding);
    }

    C10_CUDA_KERNEL_LAUNCH_CHECK();
    return out;
}
"""

_depthwise_ext = load_inline(
    name="depthwise_square_register_tiles_v1",
    cpp_sources=_cpp_source,
    cuda_sources=_cuda_source,
    functions=["depthwise_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 conv.weight.dtype == torch.float32
            and x.is_contiguous()
            and conv.weight.is_contiguous()
            and not torch.is_autocast_enabled()
            and not (
                torch.is_grad_enabled()
                and (
                    x.requires_grad
                    or conv.weight.requires_grad
                    or (conv.bias is not None and conv.bias.requires_grad)
                )
            )
        ):
            return _depthwise_ext.depthwise_cuda(
                x, conv.weight, conv.bias, conv.stride[0], conv.padding[0]
            )
        return conv(x)
```