```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 <cub/cub.cuh>
#include <cuda_runtime.h>

constexpr int THREADS = 256;
constexpr int ITEMS = 8;
constexpr int TILE = THREADS * ITEMS;

__global__ void tile_sums(
    const float* __restrict__ x,
    float* __restrict__ sums,
    int64_t n,
    int tiles,
    int64_t tasks)
{
    using Reduce = cub::BlockReduce<float, THREADS>;
    __shared__ typename Reduce::TempStorage storage;

    for (int64_t task = blockIdx.x; task < tasks; task += gridDim.x) {
        const int64_t row = task / tiles;
        const int tile = task % tiles;
        const int64_t start = int64_t(tile) * TILE;

        float sum = 0.0f;
        #pragma unroll
        for (int i = 0; i < ITEMS; ++i) {
            const int64_t col = start + threadIdx.x + i * THREADS;
            if (col < n)
                sum += x[row * n + col];
        }

        const float total = Reduce(storage).Sum(sum);
        if (threadIdx.x == 0)
            sums[task] = total;
        __syncthreads();
    }
}

__global__ void scan_tile_sums(
    const float* __restrict__ sums,
    float* __restrict__ offsets,
    int64_t rows,
    int tiles)
{
    const int lane = threadIdx.x & 31;
    const int warp = threadIdx.x >> 5;
    const int warps_per_block = blockDim.x >> 5;

    for (int64_t row = int64_t(blockIdx.x) * warps_per_block + warp;
         row < rows;
         row += int64_t(gridDim.x) * warps_per_block)
    {
        float carry = 0.0f;
        for (int base = 0; base < tiles; base += 32) {
            const int tile = base + lane;
            float value = tile < tiles ? sums[row * tiles + tile] : 0.0f;

            #pragma unroll
            for (int delta = 1; delta < 32; delta <<= 1) {
                const float other =
                    __shfl_up_sync(0xffffffffu, value, delta);
                if (lane >= delta)
                    value += other;
            }

            const float previous =
                __shfl_up_sync(0xffffffffu, value, 1);
            if (tile < tiles)
                offsets[row * tiles + tile] =
                    carry + (lane == 0 ? 0.0f : previous);

            carry += __shfl_sync(0xffffffffu, value, 31);
        }
    }
}

__global__ void scan_tiles(
    const float* __restrict__ x,
    const float* __restrict__ offsets,
    float* __restrict__ out,
    int64_t n,
    int tiles,
    int64_t tasks)
{
    using Load = cub::BlockLoad<
        float, THREADS, ITEMS, cub::BLOCK_LOAD_WARP_TRANSPOSE>;
    using Scan = cub::BlockScan<float, THREADS>;
    using Store = cub::BlockStore<
        float, THREADS, ITEMS, cub::BLOCK_STORE_WARP_TRANSPOSE>;

    __shared__ union {
        typename Load::TempStorage load;
        typename Scan::TempStorage scan;
        typename Store::TempStorage store;
    } storage;

    for (int64_t task = blockIdx.x; task < tasks; task += gridDim.x) {
        const int64_t row = task / tiles;
        const int tile = task % tiles;
        const int64_t start = int64_t(tile) * TILE;
        const int valid = static_cast<int>(
            n - start < TILE ? n - start : TILE);

        float values[ITEMS];
        Load(storage.load).Load(x + row * n + start, values, valid, 0.0f);
        __syncthreads();

        Scan(storage.scan).InclusiveSum(values, values);
        __syncthreads();

        const float prefix = offsets[task];
        #pragma unroll
        for (int i = 0; i < ITEMS; ++i)
            values[i] += prefix;

        Store(storage.store).Store(
            out + row * (n + 1) + start + 1, values, valid);

        if (tile == 0 && threadIdx.x == 0)
            out[row * (n + 1)] = 0.0f;
        __syncthreads();
    }
}

torch::Tensor exclusive_scan_cuda(torch::Tensor x)
{
    TORCH_CHECK(x.is_cuda(), "Expected a CUDA tensor");
    TORCH_CHECK(x.scalar_type() == torch::kFloat32, "Expected float32");
    TORCH_CHECK(x.dim() == 2 && x.is_contiguous(),
                "Expected a contiguous matrix");
    TORCH_CHECK(x.size(0) > 0 && x.size(1) > 0,
                "Expected nonempty dimensions");

    const c10::cuda::CUDAGuard guard(x.device());
    const int64_t rows = x.size(0) - 1;
    const int64_t n = x.size(1);

    // The original [:-1] removes the final row, even when dim == 1.
    auto out = torch::empty({rows, n + 1}, x.options());
    if (rows == 0)
        return out;

    const int tiles = static_cast<int>((n + TILE - 1) / TILE);
    const int64_t tasks = rows * tiles;
    auto sums = torch::empty({rows, tiles}, x.options());
    auto offsets = torch::empty_like(sums);

    const int blocks = static_cast<int>(tasks < 65535 ? tasks : 65535);
    const int64_t requested_prefix_blocks = (rows + 3) / 4;
    const int prefix_blocks = static_cast<int>(
        requested_prefix_blocks < 65535 ? requested_prefix_blocks : 65535);
    cudaStream_t stream = at::cuda::getCurrentCUDAStream();

    tile_sums<<<blocks, THREADS, 0, stream>>>(
        x.data_ptr<float>(), sums.data_ptr<float>(), n, tiles, tasks);
    C10_CUDA_KERNEL_LAUNCH_CHECK();

    scan_tile_sums<<<prefix_blocks, 128, 0, stream>>>(
        sums.data_ptr<float>(), offsets.data_ptr<float>(), rows, tiles);
    C10_CUDA_KERNEL_LAUNCH_CHECK();

    scan_tiles<<<blocks, THREADS, 0, stream>>>(
        x.data_ptr<float>(), offsets.data_ptr<float>(),
        out.data_ptr<float>(), n, tiles, tasks);
    C10_CUDA_KERNEL_LAUNCH_CHECK();

    return out;
}
"""

_extension = load_inline(
    name="exclusive_cumsum_tiled_cuda_v1",
    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.shape[0] > 0
            and x.shape[1] > 0
            and not x.requires_grad
        ):
            return _extension.exclusive_scan_cuda(x)

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