```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_vector_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__ float4 max4(float4 a, float4 b) {
    return make_float4(
        pool_max(a.x, b.x),
        pool_max(a.y, b.y),
        pool_max(a.z, b.z),
        pool_max(a.w, b.w));
}

__device__ __forceinline__ float4 reduce_y(
    const float* __restrict__ base,
    int z,
    int y0,
    int column)
{
    const float4 negative_infinity = make_float4(
        -CUDART_INF_F, -CUDART_INF_F,
        -CUDART_INF_F, -CUDART_INF_F);

    if (static_cast<unsigned int>(z) >= 128)
        return negative_infinity;

    const float* row = base + z * 16384 + column;
    float4 result = negative_infinity;

    if (y0 >= 0)
        result = *reinterpret_cast<const float4*>(row + y0 * 128);

    result = max4(
        result,
        *reinterpret_cast<const float4*>(row + (y0 + 3) * 128));
    result = max4(
        result,
        *reinterpret_cast<const float4*>(row + (y0 + 6) * 128));

    return result;
}

__global__ void maxpool3d_vector_kernel(
    const float* __restrict__ input,
    float2* __restrict__ output,
    int64_t rows)
{
    const int lane = threadIdx.x;
    const int64_t row = int64_t(blockIdx.x) * 8 + threadIdx.y;

    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);

    float4 previous = reduce_y(base, first_z, y0, lane * 4);

    #pragma unroll
    for (int iteration = 0; iteration < 3; ++iteration) {
        const int od = first_od + iteration * 3;

        if (od < 62) {
            const int z0 = first_z + iteration * 6;
            const float4 middle = reduce_y(base, z0 + 3, y0, lane * 4);
            const float4 next = reduce_y(base, z0 + 6, y0, lane * 4);
            const float4 value = max4(max4(previous, middle), next);
            previous = next;

            // A warp covers all 128 input columns and 62 output columns.
            float left = __shfl_up_sync(0xffffffffu, value.w, 1);
            const float next_x =
                __shfl_down_sync(0xffffffffu, value.x, 1);
            const float next_y =
                __shfl_down_sync(0xffffffffu, value.y, 1);
            const float next_w =
                __shfl_down_sync(0xffffffffu, value.w, 1);

            if (lane == 0)
                left = -CUDART_INF_F;

            if (lane < 31) {
                const float2 result = make_float2(
                    pool_max(pool_max(left, value.z), next_y),
                    pool_max(pool_max(value.y, next_x), next_w));

                const int64_t offset =
                    ((plane * 62 + od) * 62 + oh) * 31 + lane;
                output[offset] = result;
            }
        }
    }
}

torch::Tensor maxpool3d_vector_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 4D or 5D input");
    TORCH_CHECK(x.storage_offset() % 4 == 0, "Expected float4 alignment");

    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) {
        const dim3 threads(32, 8);
        const dim3 blocks(static_cast<unsigned int>((rows + 7) / 8));
        const cudaStream_t stream =
            c10::cuda::getCurrentCUDAStream(x.get_device());

        maxpool3d_vector_kernel<<<blocks, threads, 0, stream>>>(
            x.data_ptr<float>(),
            reinterpret_cast<float2*>(output.data_ptr<float>()),
            rows);
        C10_CUDA_KERNEL_LAUNCH_CHECK();
    }

    return output;
}
"""

_extension = load_inline(
    name="maxpool3d_dilated_128_vector4_v6",
    cpp_sources=_cpp_source,
    cuda_sources=_cuda_source,
    functions=["maxpool3d_vector_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() % 4 == 0
            and not (torch.is_grad_enabled() and x.requires_grad)
        ):
            return _extension.maxpool3d_vector_cuda(x)
        return self.maxpool(x)
```