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

# Fused split-K matrix-vector multiplication (C = A @ B) CUDA kernel.
#
# Workload: A (M, K) = (2048, 1048576), B (K, 1) -> C (M, 1).
# Purely DRAM-bandwidth bound (~8 GiB of A streamed exactly once), so the
# whole design is about keeping the memory controllers saturated:
#   * Split-K: the K dimension is partitioned across kSplitK blocks per row
#     group, giving 4096 independent blocks. This keeps every SM saturated
#     and smooths load balance (no long tail wave).
#   * float4 (128-bit) fully-coalesced loads.
#   * Main loop unrolled 4x (4 float4 column chunks per thread per
#     iteration): 16 outstanding 128-bit A loads per thread give deep
#     memory-level parallelism, and with slice4 a multiple of 1024 the
#     unrolled body runs remainder-free (one loop branch per 16 KB per
#     block-row). Address offsets collapse into LDG immediate offsets.
#   * __ldcs (evict-first streaming) for A: each element is read once, so it
#     must not displace B in L2.
#   * __ldg (read-only) for B: the 4 MiB vector stays resident in L2 and is
#     shared by all blocks; each load feeds kRowsPerBlock dot products.
#   * Warp-shuffle + shared-memory block reduction, finished with a single
#     red.global.add (atomicAdd) per row slice - no global synchronization
#     and no second reduction kernel.

matvec_cuda_source = r"""
#include <torch/extension.h>
#include <ATen/cuda/CUDAContext.h>
#include <c10/cuda/CUDAException.h>
#include <cuda_runtime.h>

namespace {

constexpr int kBlockSize    = 256;
constexpr int kNumWarps     = kBlockSize / 32;
constexpr int kRowsPerBlock = 4;
constexpr int kSplitK       = 8;

template <int R>
__global__ void __launch_bounds__(kBlockSize)
matvec_splitk_kernel(const float4* __restrict__ A4,
                     const float4* __restrict__ B4,
                     float* __restrict__ C,
                     const int M,
                     const int K4,
                     const int slice4) {
    const int row0 = blockIdx.x * R;
    const int i0   = blockIdx.y * slice4;
    const int i1   = min(K4, i0 + slice4);

    // One base pointer per row handled by this block (clamped on tail blocks;
    // those recomputed rows are valid reads and are masked on write).
    const float4* a_rows[R];
    #pragma unroll
    for (int r = 0; r < R; ++r) {
        int row = row0 + r;
        if (row >= M) row = M - 1;
        a_rows[r] = A4 + static_cast<size_t>(row) * K4;
    }

    float4 acc[R];
    #pragma unroll
    for (int r = 0; r < R; ++r) acc[r] = make_float4(0.f, 0.f, 0.f, 0.f);

    const int stride = kBlockSize;
    int i = i0 + threadIdx.x;

    // Main loop: four float4 column chunks per iteration. When slice4 is a
    // multiple of 4*stride (true for this model) this loop is exact and the
    // remainder loop below never runs.
    for (; i + 3 * stride < i1; i += 4 * stride) {
        const float4 b0 = __ldg(B4 + i);
        const float4 b1 = __ldg(B4 + i + stride);
        const float4 b2 = __ldg(B4 + i + 2 * stride);
        const float4 b3 = __ldg(B4 + i + 3 * stride);
        #pragma unroll
        for (int r = 0; r < R; ++r) {
            const float4 a0 = __ldcs(a_rows[r] + i);
            const float4 a1 = __ldcs(a_rows[r] + i + stride);
            const float4 a2 = __ldcs(a_rows[r] + i + 2 * stride);
            const float4 a3 = __ldcs(a_rows[r] + i + 3 * stride);
            acc[r].x += a0.x * b0.x + a1.x * b1.x + a2.x * b2.x + a3.x * b3.x;
            acc[r].y += a0.y * b0.y + a1.y * b1.y + a2.y * b2.y + a3.y * b3.y;
            acc[r].z += a0.z * b0.z + a1.z * b1.z + a2.z * b2.z + a3.z * b3.z;
            acc[r].w += a0.w * b0.w + a1.w * b1.w + a2.w * b2.w + a3.w * b3.w;
        }
    }
    // Slice remainder.
    for (; i < i1; i += stride) {
        const float4 b = __ldg(B4 + i);
        #pragma unroll
        for (int r = 0; r < R; ++r) {
            const float4 a = __ldcs(a_rows[r] + i);
            acc[r].x += a.x * b.x;
            acc[r].y += a.y * b.y;
            acc[r].z += a.z * b.z;
            acc[r].w += a.w * b.w;
        }
    }

    // Collapse float4 lanes, then reduce across the block.
    float sum[R];
    #pragma unroll
    for (int r = 0; r < R; ++r) {
        sum[r] = (acc[r].x + acc[r].y) + (acc[r].z + acc[r].w);
    }

    #pragma unroll
    for (int r = 0; r < R; ++r) {
        #pragma unroll
        for (int off = 16; off > 0; off >>= 1) {
            sum[r] += __shfl_down_sync(0xffffffffu, sum[r], off);
        }
    }

    __shared__ float warp_sums[R][kNumWarps];
    const int lane = threadIdx.x & 31;
    const int wid  = threadIdx.x >> 5;

    if (lane == 0) {
        #pragma unroll
        for (int r = 0; r < R; ++r) warp_sums[r][wid] = sum[r];
    }
    __syncthreads();

    if (wid == 0) {
        #pragma unroll
        for (int r = 0; r < R; ++r) {
            float s = (lane < kNumWarps) ? warp_sums[r][lane] : 0.f;
            #pragma unroll
            for (int off = kNumWarps / 2; off > 0; off >>= 1) {
                s += __shfl_down_sync(0xffffffffu, s, off);
            }
            if (lane == 0) {
                const int row = row0 + r;
                if (row < M) atomicAdd(&C[row], s);
            }
        }
    }
}

}  // namespace

torch::Tensor matvec_cuda(torch::Tensor A, torch::Tensor B) {
    TORCH_CHECK(A.is_cuda() && B.is_cuda(), "A and B must be CUDA tensors");
    TORCH_CHECK(A.scalar_type() == at::kFloat && B.scalar_type() == at::kFloat,
                "A and B must be float32");
    TORCH_CHECK(A.dim() == 2 && B.dim() == 2 && B.size(1) == 1,
                "Expected shapes A (M, K) and B (K, 1)");
    TORCH_CHECK(A.size(1) == B.size(0), "Inner dimensions must match");

    const int64_t M = A.size(0);
    const int64_t K = A.size(1);
    auto C = torch::empty({M, 1}, A.options());

    if (M == 0 || K == 0) {
        return C.zero_();
    }
    if ((K & 3) != 0) {
        // Alignment fallback (never triggered by this model's shapes).
        C.copy_(torch::matmul(A, B));
        return C;
    }

    const int K4 = static_cast<int>(K >> 2);

    // Split the K dimension so that roughly kSplitK blocks cooperate on each
    // row group. Keep slices a multiple of 4 * blockDim float4 chunks so the
    // unrolled main loop is exact and every warp access stays 128B aligned.
    int slice4 = (K4 + kSplitK - 1) / kSplitK;
    const int align = kBlockSize * 4;
    slice4 = ((slice4 + align - 1) / align) * align;
    const int num_slices = (K4 + slice4 - 1) / slice4;

    const dim3 grid(
        static_cast<unsigned int>((M + kRowsPerBlock - 1) / kRowsPerBlock),
        static_cast<unsigned int>(num_slices));

    cudaStream_t stream = at::cuda::getCurrentCUDAStream();
    C10_CUDA_CHECK(cudaMemsetAsync(C.data_ptr<float>(), 0,
                                   static_cast<size_t>(M) * sizeof(float),
                                   stream));

    matvec_splitk_kernel<kRowsPerBlock><<<grid, kBlockSize, 0, stream>>>(
        reinterpret_cast<const float4*>(A.data_ptr<float>()),
        reinterpret_cast<const float4*>(B.data_ptr<float>()),
        C.data_ptr<float>(),
        static_cast<int>(M),
        K4,
        slice4);
    C10_CUDA_KERNEL_LAUNCH_CHECK();

    return C;
}
"""

matvec_cpp_source = "torch::Tensor matvec_cuda(torch::Tensor A, torch::Tensor B);"

matvec_ext = load_inline(
    name="matvec_splitk_u4_ext",
    cpp_sources=matvec_cpp_source,
    cuda_sources=matvec_cuda_source,
    functions=["matvec_cuda"],
    verbose=False,
    extra_cflags=["-O3"],
    extra_cuda_cflags=["-O3"],
)


class ModelNew(nn.Module):
    """
    Optimized matrix-vector multiplication (C = A @ B) backed by a custom
    split-K, bandwidth-saturating CUDA kernel (4x-unrolled float4 streaming
    loads, single-kernel atomic reduction). Falls back to torch.matmul for
    any unsupported input configuration.
    """

    def __init__(self):
        super(ModelNew, self).__init__()
        self._fn = matvec_ext.matvec_cuda

    def forward(self, A: torch.Tensor, B: torch.Tensor) -> torch.Tensor:
        if (
            A.is_cuda
            and B.is_cuda
            and A.dtype == torch.float32
            and B.dtype == torch.float32
            and A.dim() == 2
            and B.dim() == 2
            and B.size(1) == 1
            and A.size(1) == B.size(0)
        ):
            return self._fn(A.contiguous(), B.contiguous())
        return torch.matmul(A, B)
```