```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 maxpool3d_optimized(torch::Tensor x);
"""

_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 <cmath>
#include <cstdint>

__device__ __forceinline__ float pool_max(float a, float b) {
    return (isnan(b) || b > a) ? b : a;
}

// Specialization: spatial shape 128^3, kernel=3, stride=2,
// padding=1, dilation=3. Output spatial shape is 62^3.
//
// Each warp computes 32 consecutive output columns. The first and
// third width taps are separated by three lanes, allowing their
// depth/height reductions to be shared through warp shuffles.
__global__ void maxpool3d_kernel(
    const float* __restrict__ input,
    float* __restrict__ output,
    int64_t total_rows
) {
    constexpr int OUT = 62;
    constexpr int IN = 128;
    constexpr int WARPS_PER_BLOCK = 8;

    const int lane = threadIdx.x & 31;
    const int64_t job =
        static_cast<int64_t>(blockIdx.x) * WARPS_PER_BLOCK +
        (threadIdx.x >> 5);

    const int64_t row = job >> 1;
    if (row >= total_rows) {
        return;
    }

    const int tile = static_cast<int>(job & 1);
    const int ow = tile * 32 + lane;

    const int oh = static_cast<int>(row % OUT);
    const int od = static_cast<int>((row / OUT) % OUT);
    const int64_t nc = row / (OUT * OUT);

    const int iz0 = od * 2 - 1;
    const int iy0 = oh * 2 - 1;
    const int ix0 = ow * 2 - 1;

    const float* volume = input + nc * IN * IN * IN;

    float first = -CUDART_INF_F;
    float middle = -CUDART_INF_F;
    float halo = -CUDART_INF_F;

    #pragma unroll
    for (int kz = 0; kz < 3; ++kz) {
        const int iz = iz0 + kz * 3;

        #pragma unroll
        for (int ky = 0; ky < 3; ++ky) {
            const int iy = iy0 + ky * 3;

            if (static_cast<unsigned>(iz) < IN &&
                static_cast<unsigned>(iy) < IN) {
                const int offset = iz * IN * IN + iy * IN;

                if (static_cast<unsigned>(ix0) < IN) {
                    first = pool_max(first, volume[offset + ix0]);
                }

                if (static_cast<unsigned>(ix0 + 3) < IN) {
                    middle = pool_max(middle, volume[offset + ix0 + 3]);
                }

                // These lanes cannot obtain the third tap from
                // another lane in their warp.
                if (lane >= 29 &&
                    static_cast<unsigned>(ix0 + 6) < IN) {
                    halo = pool_max(halo, volume[offset + ix0 + 6]);
                }
            }
        }
    }

    float third = __shfl_down_sync(0xffffffffu, first, 3);
    if (lane >= 29) {
        third = halo;
    }

    if (ow < OUT) {
        output[row * OUT + ow] =
            pool_max(pool_max(first, middle), third);
    }
}

torch::Tensor maxpool3d_optimized(torch::Tensor x) {
    TORCH_CHECK(x.is_cuda(), "Expected a CUDA tensor");
    TORCH_CHECK(x.scalar_type() == at::kFloat, "Expected float32");
    TORCH_CHECK(x.is_contiguous(), "Expected contiguous input");
    TORCH_CHECK(
        x.dim() == 5 &&
        x.size(2) == 128 &&
        x.size(3) == 128 &&
        x.size(4) == 128,
        "Expected spatial dimensions (128, 128, 128)"
    );

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

    auto output = torch::empty(
        {x.size(0), x.size(1), 62, 62, 62},
        x.options()
    );

    const int64_t rows = x.size(0) * x.size(1) * 62 * 62;
    if (rows == 0) {
        return output;
    }

    constexpr int threads = 256;
    constexpr int warps_per_block = threads / 32;
    const int64_t jobs = rows * 2;
    const unsigned int blocks =
        static_cast<unsigned int>(
            (jobs + warps_per_block - 1) / warps_per_block
        );

    maxpool3d_kernel<<<
        blocks, threads, 0, at::cuda::getCurrentCUDAStream()
    >>>(
        x.data_ptr<float>(),
        output.data_ptr<float>(),
        rows
    );

    C10_CUDA_KERNEL_LAUNCH_CHECK();
    return output;
}
"""

_maxpool_extension = load_inline(
    name="maxpool3d_dilated_warp_v1",
    cpp_sources=_cpp_source,
    cuda_sources=_cuda_source,
    functions=["maxpool3d_optimized"],
    extra_cflags=["-O3"],
    extra_cuda_cflags=["-O3"],
    verbose=False,
)


def _triple(value):
    if isinstance(value, int):
        return (value, value, value)
    return tuple(value)


class ModelNew(nn.Module):
    def __init__(
        self,
        kernel_size: int,
        stride: int = None,
        padding: int = 0,
        dilation: int = 1,
        return_indices: bool = False,
        ceil_mode: bool = False,
    ):
        super().__init__()

        self.maxpool = nn.MaxPool3d(
            kernel_size=kernel_size,
            stride=stride,
            padding=padding,
            dilation=dilation,
            return_indices=return_indices,
            ceil_mode=ceil_mode,
        )

        actual_stride = kernel_size if stride is None else stride
        self._specialized = (
            _triple(kernel_size) == (3, 3, 3)
            and _triple(actual_stride) == (2, 2, 2)
            and _triple(padding) == (1, 1, 1)
            and _triple(dilation) == (3, 3, 3)
            and not return_indices
            and not ceil_mode
        )

    def forward(self, x: torch.Tensor):
        if (
            self._specialized
            and x.is_cuda
            and x.dtype == torch.float32
            and x.ndim == 5
            and x.shape[2:] == (128, 128, 128)
            and x.shape[1] > 0
            and x.is_contiguous()
            and not (torch.is_grad_enabled() and x.requires_grad)
        ):
            return _maxpool_extension.maxpool3d_optimized(x)

        return self.maxpool(x)
```