```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 exclusive_scan_cuda(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>

constexpr int THREADS = 1024;
constexpr int WARPS = THREADS / 32;
constexpr unsigned MASK = 0xffffffffu;

__device__ __forceinline__ float warp_scan(float value) {
    const int lane = threadIdx.x & 31;
    #pragma unroll
    for (int delta = 1; delta < 32; delta <<= 1) {
        const float other = __shfl_up_sync(MASK, value, delta);
        if (lane >= delta)
            value += other;
    }
    return value;
}

__device__ __forceinline__ float sum4(const float4 v) {
    return ((v.x + v.y) + v.z) + v.w;
}

template<int ROUNDS, int PAIRS>
__global__ __launch_bounds__(THREADS, 1)
void row_scan_kernel(
    const float* __restrict__ input,
    float* __restrict__ output)
{
    constexpr int ITEMS = PAIRS * 4;
    constexpr int COLUMNS = THREADS * ROUNDS * ITEMS;
    constexpr int WARP_COLUMNS = 32 * ROUNDS * ITEMS;
    constexpr int OUTPUT_STRIDE = COLUMNS + 4;

    const int lane = threadIdx.x & 31;
    const int warp = threadIdx.x >> 5;
    const int64_t row = blockIdx.x;

    const float4* src = reinterpret_cast<const float4*>(
        input + row * COLUMNS + warp * WARP_COLUMNS);

    float4 values[ROUNDS][PAIRS];
    float prefixes[ROUNDS];
    float round_totals[ROUNDS];
    float warp_total = 0.0f;

    #pragma unroll
    for (int r = 0; r < ROUNDS; ++r) {
        float total = 0.0f;

        #pragma unroll
        for (int p = 0; p < PAIRS; ++p) {
            const float4 v = src[(r * 32 + lane) * PAIRS + p];
            values[r][p] = v;
            total += sum4(v);
        }

        const float inclusive = warp_scan(total);
        const float previous = __shfl_up_sync(MASK, inclusive, 1);
        prefixes[r] = lane == 0 ? 0.0f : previous;
        round_totals[r] = __shfl_sync(MASK, inclusive, 31);
        warp_total += round_totals[r];
    }

    __shared__ float warp_offsets[WARPS];
    if (lane == 0)
        warp_offsets[warp] = warp_total;
    __syncthreads();

    float* output_row = output + row * OUTPUT_STRIDE;

    if (warp == 0) {
        const float inclusive = warp_scan(warp_offsets[lane]);
        const float previous = __shfl_up_sync(MASK, inclusive, 1);
        warp_offsets[lane] = lane == 0 ? 0.0f : previous;
        if (lane == 31)
            output_row[COLUMNS] = inclusive;
    }
    __syncthreads();

    float carry = warp_offsets[warp];
    float4* dst = reinterpret_cast<float4*>(
        output_row + warp * WARP_COLUMNS);

    #pragma unroll
    for (int r = 0; r < ROUNDS; ++r) {
        float prefix = carry + prefixes[r];

        #pragma unroll
        for (int p = 0; p < PAIRS; ++p) {
            const float4 v = values[r][p];
            const float s2 = v.x + v.y;
            const float s3 = s2 + v.z;

            float4 result;
            result.x = prefix;
            result.y = prefix + v.x;
            result.z = prefix + s2;
            result.w = prefix + s3;

            dst[(r * 32 + lane) * PAIRS + p] = result;
            prefix += s3 + v.w;
        }

        carry += round_totals[r];
    }
}

torch::Tensor exclusive_scan_cuda(torch::Tensor x) {
    const c10::cuda::CUDAGuard guard(x.device());

    const int64_t rows = x.size(0) - 1;
    const int64_t columns = x.size(1);

    auto output = torch::empty_strided(
        {rows, columns + 1},
        {columns + 4, 1},
        x.options());

    if (rows == 0)
        return output;

    const auto stream = at::cuda::getCurrentCUDAStream();
    const dim3 grid(static_cast<unsigned int>(rows));

    switch (columns) {
        case 4096:
            row_scan_kernel<1, 1><<<grid, THREADS, 0, stream>>>(
                x.data_ptr<float>(), output.data_ptr<float>());
            break;
        case 8192:
            row_scan_kernel<1, 2><<<grid, THREADS, 0, stream>>>(
                x.data_ptr<float>(), output.data_ptr<float>());
            break;
        case 16384:
            row_scan_kernel<2, 2><<<grid, THREADS, 0, stream>>>(
                x.data_ptr<float>(), output.data_ptr<float>());
            break;
        case 32768:
            row_scan_kernel<4, 2><<<grid, THREADS, 0, stream>>>(
                x.data_ptr<float>(), output.data_ptr<float>());
            break;
        default:
            TORCH_CHECK(false, "Unsupported row length");
    }

    C10_CUDA_KERNEL_LAUNCH_CHECK();
    return output;
}
"""

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


class ModelNew(nn.Module):
    def __init__(self, dim):
        super().__init__()
        self.dim = dim

    def forward(self, x):
        if (
            x.is_cuda
            and x.dtype == torch.float32
            and x.ndim == 2
            and self.dim in (1, -1)
            and x.is_contiguous()
            and x.data_ptr() % 16 == 0
            and 0 < x.size(0) <= 2147483647
            and x.size(1) in (4096, 8192, 16384, 32768)
            and not (torch.is_grad_enabled() and x.requires_grad)
        ):
            # Preserve the reference's slicing along dimension zero.
            return _extension.exclusive_scan_cuda(x)

        shifted = torch.cat(
            (
                torch.zeros_like(x.select(self.dim, 0).unsqueeze(self.dim)),
                x,
            ),
            dim=self.dim,
        )[:-1]
        return torch.cumsum(shifted, dim=self.dim)
```