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


_cpp_source = r"""
#include <torch/extension.h>
torch::Tensor maxpool3d_fast_cuda(torch::Tensor x);
"""

_cuda_source = r"""
#include <torch/extension.h>
#include <c10/cuda/CUDAGuard.h>
#include <c10/cuda/CUDAStream.h>
#include <c10/cuda/CUDAException.h>
#include <cuda_runtime.h>
#include <math_constants.h>
#include <stdint.h>

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

__device__ __forceinline__ float2 reduce_y(
    const float* __restrict__ base,
    int z,
    int y0,
    int column)
{
    float2 result = make_float2(-CUDART_INF_F, -CUDART_INF_F);

    #pragma unroll
    for (int ky = 0; ky < 3; ++ky) {
        const int y = y0 + 3 * ky;
        if (z >= 0 && z < 128 && y >= 0) {
            const float2 value = *reinterpret_cast<const float2*>(
                base + z * 16384 + y * 128 + column);
            result.x = pool_max(result.x, value.x);
            result.y = pool_max(result.y, value.y);
        }
    }
    return result;
}

__device__ __forceinline__ float reduce_yz_column(
    const float* __restrict__ base,
    int z0,
    int y0,
    int column)
{
    float result = -CUDART_INF_F;

    #pragma unroll
    for (int kz = 0; kz < 3; ++kz) {
        const int z = z0 + 3 * kz;
        if (z >= 0) {
            #pragma unroll
            for (int ky = 0; ky < 3; ++ky) {
                const int y = y0 + 3 * ky;
                if (y >= 0) {
                    result = pool_max(
                        result, base[z * 16384 + y * 128 + column]);
                }
            }
        }
    }
    return result;
}

__global__ void maxpool3d_128_reuse_kernel(
    const float* __restrict__ input,
    float* __restrict__ output,
    int64_t rows)
{
    const int ow = threadIdx.x;
    const int lane = ow & 31;
    const int64_t row = int64_t(blockIdx.x) * 4 + threadIdx.y;

    // Each complete warp either participates or exits.
    if (row >= rows)
        return;

    const int oh = row % 62;
    const int64_t grouped_row = row / 62;
    const int depth_group = grouped_row % 21;
    const int64_t plane = grouped_row / 21;
    const int first_od = (depth_group / 3) * 9 + depth_group % 3;
    const int first_z = first_od * 2 - 1;
    const int y0 = oh * 2 - 1;

    const float* base = input + plane * int64_t(128 * 128 * 128);

    // Adjacent threads read adjacent float2 values. Three output depths
    // separated by three share the boundary depth of their pooling windows.
    float2 previous = reduce_y(base, first_z, y0, 2 * ow);

    #pragma unroll
    for (int iteration = 0; iteration < 3; ++iteration) {
        const int od = first_od + 3 * iteration;
        if (od < 62) {
            const int z0 = first_z + 6 * iteration;
            const float2 middle = reduce_y(base, z0 + 3, y0, 2 * ow);
            const float2 next = reduce_y(base, z0 + 6, y0, 2 * ow);

            const float even = pool_max(
                pool_max(previous.x, middle.x), next.x);
            const float odd = pool_max(
                pool_max(previous.y, middle.y), next.y);
            previous = next;

            // The width window is [2*ow-1, 2*ow+2, 2*ow+5].
            float left = __shfl_up_sync(0xffffffffu, odd, 1);
            float center = __shfl_down_sync(0xffffffffu, even, 1);
            float right = __shfl_down_sync(0xffffffffu, odd, 2);

            if (ow < 62) {
                if (lane == 0) {
                    left = ow == 0
                        ? -CUDART_INF_F
                        : reduce_yz_column(base, z0, y0, 2 * ow - 1);
                }
                if (lane == 31) {
                    center = reduce_yz_column(
                        base, z0, y0, 2 * ow + 2);
                }
                if (lane >= 30) {
                    right = reduce_yz_column(
                        base, z0, y0, 2 * ow + 5);
                }

                const int64_t output_offset =
                    ((plane * 62 + od) * 62 + oh) * 62 + ow;
                output[output_offset] =
                    pool_max(pool_max(left, center), right);
            }
        }
    }
}

torch::Tensor maxpool3d_fast_cuda(torch::Tensor x)
{
    TORCH_CHECK(x.is_cuda(), "Expected a CUDA tensor");
    TORCH_CHECK(x.scalar_type() == at::kFloat, "Expected float32 input");
    TORCH_CHECK(x.is_contiguous(), "Expected contiguous input");
    TORCH_CHECK(x.dim() == 4 || x.dim() == 5, "Expected a 4D or 5D input");

    const int rank = x.dim();
    TORCH_CHECK(
        x.size(rank - 3) == 128 &&
        x.size(rank - 2) == 128 &&
        x.size(rank - 1) == 128,
        "Expected spatial dimensions 128 x 128 x 128");

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

    auto shape = x.sizes().vec();
    shape[rank - 3] = 62;
    shape[rank - 2] = 62;
    shape[rank - 1] = 62;
    auto output = torch::empty(shape, x.options());

    const int64_t planes = x.numel() / (128 * 128 * 128);
    const int64_t rows = planes * 21 * 62;
    if (rows == 0)
        return output;

    const dim3 threads(64, 4);
    const dim3 blocks(static_cast<unsigned int>((rows + 3) / 4));
    const cudaStream_t stream =
        c10::cuda::getCurrentCUDAStream(x.get_device());

    maxpool3d_128_reuse_kernel<<<blocks, threads, 0, stream>>>(
        x.data_ptr<float>(), output.data_ptr<float>(), rows);
    C10_CUDA_KERNEL_LAUNCH_CHECK();

    return output;
}
"""

_extension = load_inline(
    name="maxpool3d_dilated_128_reuse_v4",
    cpp_sources=_cpp_source,
    cuda_sources=_cuda_source,
    functions=["maxpool3d_fast_cuda"],
    extra_cflags=["-O3"],
    extra_cuda_cflags=["-O3"],
    verbose=False,
)


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,
        )
        self._specialized = (
            tuple(_triple(kernel_size)) == (3, 3, 3)
            and tuple(_triple(kernel_size if stride is None else stride))
            == (2, 2, 2)
            and tuple(_triple(padding)) == (1, 1, 1)
            and tuple(_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 in (4, 5)
            and x.shape[-3:] == (128, 128, 128)
            and x.shape[-4] > 0
            and x.is_contiguous()
            and x.storage_offset() % 2 == 0
            and not (torch.is_grad_enabled() and x.requires_grad)
        ):
            return _extension.maxpool3d_fast_cuda(x)
        return self.maxpool(x)
```