# Optimized Model with Custom CUDA Kernel

The original model has a subtle behavior: `torch.cat(...)[ :-1]` slices **dimension 0**, so for input `(32768, 32768)` with `dim=1`, the output is actually shape `(32767, 32769)` where `out[i,0]=0` and `out[i,j]=sum_{k<j} x[i,k]`. My kernel reproduces this exactly while being fully memory-optimal (single read of input, single write of output, ~8GB total traffic vs ~16GB+ for the baseline).

Key optimizations:
- **Single fused kernel** replaces `cat` + slice + `cumsum` (eliminates 4GB intermediate tensor).
- **One block per row**, striped coalesced loads/stores, 4 elements per thread for ILP.
- **Warp-shuffle block scan** with `double` accumulators (near-exact numerics, still memory-bound).

```python
# Copyright (c) Meta Platforms, Inc. and affiliates.
# All rights reserved.
#
# This source code is licensed under the license found in the
# LICENSE file in the root directory of this source tree.

import torch
import torch.nn as nn
from torch.utils.cpp_extension import load_inline

_cuda_source = r"""
#include <torch/extension.h>
#include <ATen/cuda/CUDAContext.h>
#include <cuda_runtime.h>
#include <cstdint>

#define FULL_MASK 0xffffffffu

// Fused exclusive-cumsum reproducing:
//   torch.cat((zeros_like(x.select(dim,0).unsqueeze(dim)), x), dim=dim)[:-1].cumsum(dim)
// for the case dim == last dimension.
//
// One thread block scans one row of length n.
//   SHIFTED == true  (ndim >= 2): out row length = n+1, out[0] = 0,
//                                 out[j] = sum_{k<j} x[k] for j = 1..n
//   SHIFTED == false (ndim == 1): out row length = n,
//                                 out[j] = sum_{k<j} x[k] for j = 0..n-1
template <int THREADS, int ITEMS, bool SHIFTED>
__global__ void exclusive_scan_rows_kernel(const float* __restrict__ x,
                                           float* __restrict__ out,
                                           const int64_t n) {
    const int64_t row = blockIdx.x;
    const int64_t out_stride = SHIFTED ? (n + 1) : n;
    const float* __restrict__ xr = x + row * n;
    float* __restrict__ outr = out + row * out_stride;

    constexpr int NWARPS = THREADS / 32;
    __shared__ double warp_sums[NWARPS];

    const int tid = threadIdx.x;
    const int lane = tid & 31;
    const int warp = tid >> 5;

    if (SHIFTED && tid == 0) outr[0] = 0.0f;

    double carry = 0.0;
    const int64_t TILE = (int64_t)THREADS * ITEMS;

    for (int64_t base = 0; base < n; base += TILE) {
        // Striped, fully coalesced loads; ITEMS elements per thread.
        double v[ITEMS];
        double tsum = 0.0;
        #pragma unroll
        for (int i = 0; i < ITEMS; ++i) {
            const int64_t idx = base + tid + (int64_t)i * THREADS;
            v[i] = (idx < n) ? (double)xr[idx] : 0.0;
            tsum += v[i];
        }

        // Inclusive scan of per-thread sums across the block (warp shuffles).
        double incl = tsum;
        #pragma unroll
        for (int d = 1; d < 32; d <<= 1) {
            const double t = __shfl_up_sync(FULL_MASK, incl, d);
            if (lane >= d) incl += t;
        }
        if (lane == 31) warp_sums[warp] = incl;
        __syncthreads();
        if (warp == 0) {
            double w = (lane < NWARPS) ? warp_sums[lane] : 0.0;
            #pragma unroll
            for (int d = 1; d < NWARPS; d <<= 1) {
                const double t = __shfl_up_sync(FULL_MASK, w, d);
                if (lane >= d) w += t;
            }
            if (lane < NWARPS) warp_sums[lane] = w;
        }
        __syncthreads();

        double run = carry + (incl - tsum) + (warp > 0 ? warp_sums[warp - 1] : 0.0);
        const double total = warp_sums[NWARPS - 1];

        // Coalesced stores.
        #pragma unroll
        for (int i = 0; i < ITEMS; ++i) {
            const int64_t idx = base + tid + (int64_t)i * THREADS;
            run += v[i];  // inclusive prefix of x[idx]
            if (idx < n) {
                if (SHIFTED) outr[idx + 1] = (float)run;
                else         outr[idx]     = (float)(run - v[i]);
            }
        }
        carry += total;
        __syncthreads();
    }
}

torch::Tensor exclusive_cumsum_cuda(torch::Tensor x, int64_t dim) {
    TORCH_CHECK(x.is_cuda(), "x must be a CUDA tensor");
    TORCH_CHECK(x.scalar_type() == at::kFloat, "x must be float32");
    auto xc = x.contiguous();
    const int64_t nd = xc.dim();
    TORCH_CHECK(nd >= 1, "x must have at least 1 dimension");
    if (dim < 0) dim += nd;
    TORCH_CHECK(dim == nd - 1, "fast path requires scanning the last dimension");

    const int64_t n = (int64_t)xc.size(dim);
    cudaStream_t stream = at::cuda::getCurrentCUDAStream();

    constexpr int THREADS = 512;
    constexpr int ITEMS = 4;

    if (nd == 1) {
        auto out = torch::empty_like(xc);
        if (n > 0) {
            exclusive_scan_rows_kernel<THREADS, ITEMS, false>
                <<<1, THREADS, 0, stream>>>(xc.data_ptr<float>(),
                                            out.data_ptr<float>(), n);
        }
        return out;
    }

    auto sizes = xc.sizes().vec();
    const int64_t size0 = (int64_t)sizes[0];
    TORCH_CHECK(size0 >= 1, "first dimension must be non-empty");
    int64_t outer = 1;
    for (int64_t i = 0; i < nd - 1; ++i) outer *= (int64_t)sizes[i];
    // The trailing [:-1] in the reference drops the last slice of dim 0.
    const int64_t rows = outer - outer / size0;

    sizes[0] -= 1;
    sizes[nd - 1] += 1;
    auto out = torch::empty(sizes, xc.options());

    if (rows > 0 && n > 0) {
        TORCH_CHECK(rows <= 2147483647LL, "too many rows");
        exclusive_scan_rows_kernel<THREADS, ITEMS, true>
            <<<(unsigned int)rows, THREADS, 0, stream>>>(
                xc.data_ptr<float>(), out.data_ptr<float>(), n);
    }
    return out;
}
"""

_cpp_source = "torch::Tensor exclusive_cumsum_cuda(torch::Tensor x, int64_t dim);"

_ext = load_inline(
    name="exclusive_cumsum_ext",
    cpp_sources=_cpp_source,
    cuda_sources=_cuda_source,
    functions=["exclusive_cumsum_cuda"],
    verbose=False,
    extra_cflags=["-O3"],
    extra_cuda_cflags=["-O3"],
)


class ModelNew(nn.Module):
    """
    Optimized exclusive cumulative sum. Uses a single fused custom CUDA kernel
    (scan along the last dimension) instead of cat + slice + cumsum, removing
    the large intermediate tensor and halving memory traffic. Falls back to the
    reference implementation for unsupported configurations.
    """

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

    def forward(self, x):
        if x.is_cuda and x.dtype == torch.float32 and x.dim() >= 1 and x.shape[0] >= 1:
            d = self.dim + x.dim() if self.dim < 0 else self.dim
            if d == x.dim() - 1:
                return _ext.exclusive_cumsum_cuda(x, self.dim)
        # Reference fallback (exact original semantics).
        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)
```