```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 splitk_matvec_cuda(torch::Tensor A, torch::Tensor B);
"""

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

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;
}

__device__ __forceinline__ float block_sum(float value) {
    __shared__ float warp_sums[THREADS / 32];

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

    value = warp_sum(value);
    if (lane == 0)
        warp_sums[warp] = value;
    __syncthreads();

    value = (threadIdx.x < THREADS / 32)
        ? warp_sums[threadIdx.x] : 0.0f;

    if (warp == 0)
        value = warp_sum(value);
    return value;
}

__global__ void matvec_partials(
    const float* __restrict__ A,
    const float* __restrict__ B,
    float* __restrict__ partials,
    int K,
    int splits
) {
    const int row = blockIdx.x / splits;
    const int split = blockIdx.x % splits;
    const int start4 = split * (CHUNK / 4);
    const int K4 = K / 4;

    const float4* __restrict__ a4 =
        reinterpret_cast<const float4*>(A + static_cast<int64_t>(row) * K);
    const float4* __restrict__ b4 =
        reinterpret_cast<const float4*>(B);

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

    #pragma unroll
    for (int offset = 0; offset < CHUNK / 4; offset += THREADS) {
        const int index = start4 + offset + threadIdx.x;
        if (index < K4) {
            const float4 av = a4[index];
            const float4 bv = b4[index];

            s0 = fmaf(av.x, bv.x, s0);
            s1 = fmaf(av.y, bv.y, s1);
            s2 = fmaf(av.z, bv.z, s2);
            s3 = fmaf(av.w, bv.w, s3);
        }
    }

    const float total = block_sum((s0 + s1) + (s2 + s3));
    if (threadIdx.x == 0)
        partials[blockIdx.x] = total;
}

__global__ void finish_matvec(
    const float* __restrict__ partials,
    float* __restrict__ output,
    int splits
) {
    const int row = blockIdx.x;
    const float* row_partials =
        partials + static_cast<int64_t>(row) * splits;

    float sum = 0.0f;
    for (int i = threadIdx.x; i < splits; i += THREADS)
        sum += row_partials[i];

    sum = block_sum(sum);
    if (threadIdx.x == 0)
        output[row] = sum;
}

torch::Tensor splitk_matvec_cuda(torch::Tensor A, torch::Tensor B) {
    TORCH_CHECK(A.is_cuda() && B.is_cuda(), "Expected CUDA tensors");
    TORCH_CHECK(A.device() == B.device(), "Device mismatch");
    TORCH_CHECK(
        A.scalar_type() == torch::kFloat32 &&
        B.scalar_type() == torch::kFloat32,
        "Expected float32 tensors"
    );
    TORCH_CHECK(A.is_contiguous() && B.is_contiguous(),
                "Expected contiguous tensors");
    TORCH_CHECK(
        A.dim() == 2 && B.dim() == 2 &&
        B.size(1) == 1 && A.size(1) == B.size(0),
        "Invalid matrix-vector dimensions"
    );

    const int64_t M = A.size(0);
    const int64_t K64 = A.size(1);
    TORCH_CHECK(K64 > 0 && K64 <= 2147483647LL && K64 % 4 == 0,
                "Unsupported reduction dimension");

    const c10::cuda::CUDAGuard device_guard(A.device());
    auto output = torch::empty({M, 1}, A.options());
    if (M == 0)
        return output;

    const int K = static_cast<int>(K64);
    const int splits = static_cast<int>((K64 + CHUNK - 1) / CHUNK);
    TORCH_CHECK(M * splits <= 2147483647LL, "Grid is too large");

    const auto stream = at::cuda::getCurrentCUDAStream();

    if (splits == 1) {
        matvec_partials<<<static_cast<unsigned int>(M), THREADS, 0, stream>>>(
            A.data_ptr<float>(), B.data_ptr<float>(),
            output.data_ptr<float>(), K, splits
        );
        C10_CUDA_KERNEL_LAUNCH_CHECK();
        return output;
    }

    auto partials = torch::empty({M, splits}, A.options());

    matvec_partials<<<
        static_cast<unsigned int>(M * splits), THREADS, 0, stream
    >>>(
        A.data_ptr<float>(), B.data_ptr<float>(),
        partials.data_ptr<float>(), K, splits
    );
    C10_CUDA_KERNEL_LAUNCH_CHECK();

    finish_matvec<<<static_cast<unsigned int>(M), THREADS, 0, stream>>>(
        partials.data_ptr<float>(), output.data_ptr<float>(), splits
    );
    C10_CUDA_KERNEL_LAUNCH_CHECK();

    return output;
}
"""

_matvec_extension = load_inline(
    name="splitk_matvec_fp32_v1",
    cpp_sources=_cpp_source,
    cuda_sources=_cuda_source,
    functions=["splitk_matvec_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[1] == 1
            and A.shape[1] == B.shape[0]
            and 0 < A.shape[1] <= 2147483647
            and A.shape[1] % 4 == 0
            and A.shape[0] * ((A.shape[1] + 16383) // 16384) <= 2147483647
            and A.is_contiguous()
            and B.is_contiguous()
            and A.data_ptr() % 16 == 0
            and B.data_ptr() % 16 == 0
            and not (torch.is_grad_enabled() and (A.requires_grad or B.requires_grad))
        ):
            return _matvec_extension.splitk_matvec_cuda(A, B)
        return torch.matmul(A, B)
```