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

# Custom split-K matrix-vector multiplication kernel (C = A @ B).
# A is (M, K) row-major, B is (K, 1), C is (M, 1).
# The operation is memory-bandwidth bound, so the kernel is designed to:
#   * use 128-bit (float4) vectorized global loads through the read-only path
#   * split each row's reduction across several thread blocks (split-K) so the
#     whole GPU stays saturated even when M is modest
#   * use 4 independent accumulators per thread for ILP
#   * reduce with warp shuffles + a tiny shared-memory second stage
#   * combine the split-K partials with atomicAdd into a zero-initialized output
_matvec_cuda_src = r"""
#include <torch/extension.h>
#include <ATen/cuda/CUDAContext.h>
#include <cuda_runtime.h>

#define MATVEC_THREADS 256

__global__ void matvec_splitk_kernel(
    const float4* __restrict__ A4,
    const float4* __restrict__ B4,
    const float*  __restrict__ As,
    const float*  __restrict__ Bs,
    float* __restrict__ C,
    const int K4,
    const int K,
    const int chunk)
{
    const int row  = blockIdx.x;
    const int part = blockIdx.y;

    const int start = part * chunk;
    int end = start + chunk;
    if (end > K4) end = K4;

    const float4* a = A4 + (size_t)row * (size_t)K4;

    const int stride = blockDim.x;
    int i = start + threadIdx.x;

    float s0 = 0.f, s1 = 0.f, s2 = 0.f, s3 = 0.f;

    // Main vectorized loop, unrolled by 4 for memory-level parallelism.
    for (; i + 3 * stride < end; i += 4 * stride) {
        float4 a0 = __ldg(a + i);
        float4 a1 = __ldg(a + i + stride);
        float4 a2 = __ldg(a + i + 2 * stride);
        float4 a3 = __ldg(a + i + 3 * stride);
        float4 b0 = __ldg(B4 + i);
        float4 b1 = __ldg(B4 + i + stride);
        float4 b2 = __ldg(B4 + i + 2 * stride);
        float4 b3 = __ldg(B4 + i + 3 * stride);
        s0 += a0.x * b0.x + a0.y * b0.y + a0.z * b0.z + a0.w * b0.w;
        s1 += a1.x * b1.x + a1.y * b1.y + a1.z * b1.z + a1.w * b1.w;
        s2 += a2.x * b2.x + a2.y * b2.y + a2.z * b2.z + a2.w * b2.w;
        s3 += a3.x * b3.x + a3.y * b3.y + a3.z * b3.z + a3.w * b3.w;
    }
    for (; i < end; i += stride) {
        float4 av = __ldg(a + i);
        float4 bv = __ldg(B4 + i);
        s0 += av.x * bv.x + av.y * bv.y + av.z * bv.z + av.w * bv.w;
    }

    float sum = (s0 + s1) + (s2 + s3);

    // Scalar tail when K is not divisible by 4 (handled by slice 0).
    if (part == 0) {
        const int base = K4 * 4;
        const float* arow = As + (size_t)row * (size_t)K;
        for (int j = base + threadIdx.x; j < K; j += stride) {
            sum += arow[j] * Bs[j];
        }
    }

    // Warp-level reduction.
    #pragma unroll
    for (int off = 16; off > 0; off >>= 1) {
        sum += __shfl_down_sync(0xffffffffu, sum, off);
    }

    __shared__ float warp_sums[MATVEC_THREADS / 32];
    const int lane = threadIdx.x & 31;
    const int warp = threadIdx.x >> 5;
    if (lane == 0) warp_sums[warp] = sum;
    __syncthreads();

    if (warp == 0) {
        float v = (lane < MATVEC_THREADS / 32) ? warp_sums[lane] : 0.f;
        #pragma unroll
        for (int off = MATVEC_THREADS / 64; off > 0; off >>= 1) {
            v += __shfl_down_sync(0xffffffffu, v, off);
        }
        if (lane == 0) atomicAdd(C + row, v);
    }
}

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, "A must be 2-D");
    TORCH_CHECK(B.numel() == A.size(1), "B must have A.size(1) elements");

    auto Ac = A.contiguous();
    auto Bc = B.contiguous();

    const int64_t M = Ac.size(0);
    const int64_t K = Ac.size(1);

    auto C = torch::zeros({M, 1}, Ac.options());
    if (M == 0 || K == 0) return C;

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

    // Pick a split-K factor so there are enough blocks to fill the GPU.
    int split = 1;
    while ((int64_t)split * M < 8192 && split < 16) split <<= 1;

    int chunk = 0;
    if (K4 > 0) {
        const int unit = MATVEC_THREADS * 4;
        chunk = (K4 + split - 1) / split;
        chunk = ((chunk + unit - 1) / unit) * unit;  // keep slices balanced
        split = (K4 + chunk - 1) / chunk;
    } else {
        split = 1;
        chunk = 0;
    }

    dim3 grid((unsigned)M, (unsigned)split);
    dim3 block(MATVEC_THREADS);

    auto stream = at::cuda::getCurrentCUDAStream();
    matvec_splitk_kernel<<<grid, block, 0, stream>>>(
        reinterpret_cast<const float4*>(Ac.data_ptr<float>()),
        reinterpret_cast<const float4*>(Bc.data_ptr<float>()),
        Ac.data_ptr<float>(),
        Bc.data_ptr<float>(),
        C.data_ptr<float>(),
        K4, (int)K, chunk);

    return C;
}
"""

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

_matvec_ext = load_inline(
    name="matvec_splitk_ext",
    cpp_sources=_matvec_cpp_src,
    cuda_sources=_matvec_cuda_src,
    functions=["matvec_cuda"],
    verbose=False,
    extra_cflags=["-O3"],
    extra_cuda_cflags=["-O3"],
)


class ModelNew(nn.Module):
    """
    Matrix-vector multiplication (C = A * B) accelerated with a custom
    split-K CUDA kernel that saturates memory bandwidth using float4 loads.
    """
    def __init__(self):
        super(ModelNew, self).__init__()
        self._op = _matvec_ext

    def forward(self, A: torch.Tensor, B: torch.Tensor) -> torch.Tensor:
        if (A.is_cuda and B.is_cuda and A.dim() == 2
                and A.dtype == torch.float32 and B.dtype == torch.float32):
            return self._op.matvec_cuda(A, B)
        return torch.matmul(A, B)
```