```python
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 <c10/cuda/CUDAGuard.h>
#include <c10/cuda/CUDAException.h>
#include <cuda_runtime.h>
#include <cstdint>

constexpr int THREADS = 256;
constexpr int CHUNK = 16384;

__device__ __forceinline__ float warp_sum(float value) {
    #pragma unroll
    for (int offset = 16; offset > 0; offset >>= 1)
        value += __shfl_down_sync(0xffffffff, value, offset);
    return value;
}

template<bool VECTORIZED>
__global__ void split_gemv(
    const float* __restrict__ A,
    const float* __restrict__ B,
    float* __restrict__ partial,
    int64_t K,
    int64_t splits)
{
    const int tid = threadIdx.x;
    const int64_t row = static_cast<int64_t>(blockIdx.x) / splits;
    const int64_t split = static_cast<int64_t>(blockIdx.x) % splits;
    const int64_t begin = split * CHUNK;
    const float* row_ptr = A + row * K;

    float s0 = 0.0f;
    float s1 = 0.0f;
    float s2 = 0.0f;
    float s3 = 0.0f;

    if (VECTORIZED) {
        #pragma unroll
        for (int i = 0; i < CHUNK / (THREADS * 4); ++i) {
            const int64_t k = begin + 4 * tid + i * THREADS * 4;
            if (k < K) {
                const float4 a =
                    *reinterpret_cast<const float4*>(row_ptr + k);
                const float4 b =
                    *reinterpret_cast<const float4*>(B + k);
                s0 = fmaf(a.x, b.x, s0);
                s1 = fmaf(a.y, b.y, s1);
                s2 = fmaf(a.z, b.z, s2);
                s3 = fmaf(a.w, b.w, s3);
            }
        }
    } else {
        #pragma unroll
        for (int i = 0; i < CHUNK / (THREADS * 4); ++i) {
            const int64_t k = begin + tid + i * THREADS * 4;
            if (k < K)
                s0 = fmaf(row_ptr[k], B[k], s0);
            if (k + THREADS < K)
                s1 = fmaf(row_ptr[k + THREADS], B[k + THREADS], s1);
            if (k + 2 * THREADS < K)
                s2 = fmaf(
                    row_ptr[k + 2 * THREADS], B[k + 2 * THREADS], s2);
            if (k + 3 * THREADS < K)
                s3 = fmaf(
                    row_ptr[k + 3 * THREADS], B[k + 3 * THREADS], s3);
        }
    }

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

    __shared__ float warp_sums[THREADS / 32];
    if ((tid & 31) == 0)
        warp_sums[tid >> 5] = sum;
    __syncthreads();

    if (tid < 32) {
        sum = tid < THREADS / 32 ? warp_sums[tid] : 0.0f;
        sum = warp_sum(sum);
        if (tid == 0)
            partial[row * splits + split] = sum;
    }
}

__global__ void finish_gemv(
    const float* __restrict__ partial,
    float* __restrict__ output,
    int64_t M,
    int64_t splits)
{
    const int lane = threadIdx.x & 31;
    const int64_t row =
        static_cast<int64_t>(blockIdx.x) * (THREADS / 32)
        + (threadIdx.x >> 5);

    if (row >= M)
        return;

    float sum = 0.0f;
    for (int64_t i = lane; i < splits; i += 32)
        sum += partial[row * splits + i];

    sum = warp_sum(sum);
    if (lane == 0)
        output[row] = sum;
}

torch::Tensor gemv_cuda(torch::Tensor A, torch::Tensor B) {
    TORCH_CHECK(A.is_cuda() && B.is_cuda(), "Inputs must be CUDA tensors");
    TORCH_CHECK(A.device() == B.device(), "Inputs must share a device");
    TORCH_CHECK(
        A.scalar_type() == at::kFloat && B.scalar_type() == at::kFloat,
        "Inputs must have float32 dtype");
    TORCH_CHECK(A.dim() == 2 && B.dim() == 2, "Inputs must be matrices");
    TORCH_CHECK(
        B.size(0) == A.size(1) && B.size(1) == 1,
        "Expected A[M, K] and B[K, 1]");
    TORCH_CHECK(A.is_contiguous() && B.is_contiguous(),
                "Inputs must be contiguous");

    const c10::cuda::CUDAGuard device_guard(A.device());
    const int64_t M = A.size(0);
    const int64_t K = A.size(1);

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

    const int64_t splits = (K + CHUNK - 1) / CHUNK;
    TORCH_CHECK(M * splits <= 2147483647LL, "Launch grid is too large");

    auto partial = torch::empty({M, splits}, A.options());
    auto output = torch::empty({M, 1}, A.options());
    const auto stream = at::cuda::getCurrentCUDAStream(A.get_device());

    const float* a = A.data_ptr<float>();
    const float* b = B.data_ptr<float>();
    const bool vectorized =
        K % 4 == 0 &&
        reinterpret_cast<uintptr_t>(a) % 16 == 0 &&
        reinterpret_cast<uintptr_t>(b) % 16 == 0;

    const unsigned int blocks = static_cast<unsigned int>(M * splits);
    if (vectorized) {
        split_gemv<true><<<blocks, THREADS, 0, stream>>>(
            a, b, partial.data_ptr<float>(), K, splits);
    } else {
        split_gemv<false><<<blocks, THREADS, 0, stream>>>(
            a, b, partial.data_ptr<float>(), K, splits);
    }
    C10_CUDA_KERNEL_LAUNCH_CHECK();

    const unsigned int output_blocks =
        static_cast<unsigned int>((M + THREADS / 32 - 1) / (THREADS / 32));
    finish_gemv<<<output_blocks, THREADS, 0, stream>>>(
        partial.data_ptr<float>(), output.data_ptr<float>(), M, splits);
    C10_CUDA_KERNEL_LAUNCH_CHECK();

    return output;
}
"""

_gemv = load_inline(
    name="split_vectorized_gemv_v1",
    cpp_sources="torch::Tensor gemv_cuda(torch::Tensor A, torch::Tensor B);",
    cuda_sources=_cuda_source,
    functions=["gemv_cuda"],
    extra_cflags=["-O3"],
    extra_cuda_cflags=["-O3"],
    verbose=False,
)


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

    def forward(self, A: torch.Tensor, B: torch.Tensor) -> torch.Tensor:
        if (
            A.is_cuda
            and B.is_cuda
            and A.device == B.device
            and A.dtype == torch.float32
            and B.dtype == torch.float32
            and A.ndim == 2
            and B.ndim == 2
            and B.shape[0] == A.shape[1]
            and B.shape[1] == 1
        ):
            return _gemv.gemv_cuda(A.contiguous(), B.contiguous())
        return torch.matmul(A, B)
```